diff --git a/.github/workflows/publish_docs.yml b/.github/workflows/publish_docs.yml index e7ef97bff..d6fc6bf35 100644 --- a/.github/workflows/publish_docs.yml +++ b/.github/workflows/publish_docs.yml @@ -34,10 +34,6 @@ jobs: - name: Build docs run: cd docs && uv run make html - - name: Generate LLM agent reference markdown - run: | - uv run jupyter nbconvert --to markdown docs/source/llm.ipynb --output llm.md --output-dir docs/build/html/ - - name: Setup Pages uses: actions/configure-pages@v4 diff --git a/.github/workflows/test_docs.yml b/.github/workflows/test_docs.yml index d38369c6e..ec86c253b 100644 --- a/.github/workflows/test_docs.yml +++ b/.github/workflows/test_docs.yml @@ -33,7 +33,3 @@ jobs: - name: Build docs run: | cd docs && uv run make html - - - name: Generate LLM agent reference markdown - run: | - uv run jupyter nbconvert --to markdown docs/source/llm.ipynb --output llm.md --output-dir docs/build/html/ diff --git a/docs/source/index.rst b/docs/source/index.rst index 56028b750..92aa02071 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -8,13 +8,6 @@ Table of Contents getting_started introduction named_tensor_notation - llm - -.. tip:: - - **For AI coding agents:** A standalone Markdown version of the LLM guide is - available at `llm.md `_ — point your agent there for a complete - reference on building effectful LLM applications. .. toctree:: :maxdepth: 1 diff --git a/docs/source/llm.ipynb b/docs/source/llm.ipynb deleted file mode 100644 index c738adcae..000000000 --- a/docs/source/llm.ipynb +++ /dev/null @@ -1,821 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "e7fda1b8", - "metadata": {}, - "source": [ - "# LLM Interface\n", - "The `effectful.handlers.llm` module provides a simplified LLM interface that uses algebraic effects for modularity. The module interface consists of:\n", - "\n", - "- A decorator `Template.define` which creates a prompt template from a callable. A template is an LLM-implemented function whose behavior is specified by a template string. When a template is called, an LLM is invoked to produce the specified behavior.\n", - "- A decorator `Tool.define` which exposes Python callables as tools that templates can call. Tool signatures and docstrings define the schema passed to the model.\n", - "- Structured output handling via `Encodable` (used internally by templates and tool calls) to serialize/deserialize Python types.\n", - "- LLM providers such as `LiteLLMProvider`, and reliability helpers like `RetryLLMHandler` and `ReplayLiteLLMProvider`, which can be composed with `handler(...)` to control execution." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "5aaf649f", - "metadata": {}, - "outputs": [], - "source": [ - "import base64\n", - "import dataclasses\n", - "import functools\n", - "import io\n", - "from typing import Literal\n", - "\n", - "import litellm\n", - "import pydantic\n", - "from IPython.display import HTML, display\n", - "from litellm.caching.caching import Cache\n", - "from PIL import Image\n", - "from pydantic import field_validator\n", - "from pydantic_core import PydanticCustomError\n", - "\n", - "from effectful.handlers.llm import Template, Tool\n", - "from effectful.handlers.llm.completions import (\n", - " LiteLLMProvider,\n", - " RetryLLMHandler,\n", - ")\n", - "from effectful.ops.semantics import NotHandled, handler\n", - "\n", - "provider = LiteLLMProvider()" - ] - }, - { - "cell_type": "markdown", - "id": "093243e0", - "metadata": {}, - "source": [ - "In the following sections, we walk through each of the mentioned components." - ] - }, - { - "cell_type": "markdown", - "id": "c1c639d3", - "metadata": {}, - "source": [ - "## Prompt Templates\n", - "\n", - "This template function writes (bad) poetry on a given theme. While difficult to implement in Python, an LLM can provide a reasonable implementation." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "1e832675", - "metadata": {}, - "outputs": [], - "source": [ - "@Template.define\n", - "def limerick(theme: str) -> str:\n", - " \"\"\"Write a limerick on the theme of {theme}. Do not use any tools.\"\"\"\n", - " raise NotHandled" - ] - }, - { - "cell_type": "markdown", - "id": "f2ca6919", - "metadata": {}, - "source": [ - "If we call the template with a provider interpretation installed, we get reasonable behavior. The LLM is nondeterministic by default, so calling the template twice with the same arguments gives us different results.\n", - "\n", - "Templates are regular callables, so can be converted to operations with `defop` if we want to override the LLM implementation in some cases." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "634f6533", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "In the sea where the shimmering fish \n", - "Dance around like a silvery wish,\n", - "They wiggle and glide,\n", - "With the tide, side by side,\n", - "Turning waves into their swirlish dish.\n", - "----------------------------------------\n", - "There once was a fish named Blue,\n", - "Who swam in a sea of bright hue.\n", - "With scales shining bright,\n", - "He'd dance in the light,\n", - "And none were as charming as Blue.\n" - ] - } - ], - "source": [ - "with handler(provider):\n", - " print(limerick(\"fish\"))\n", - " print(\"-\" * 40)\n", - " print(limerick(\"fish\"))" - ] - }, - { - "cell_type": "markdown", - "id": "2e59acbc", - "metadata": {}, - "source": [ - "If we want deterministic behavior, we can cache the template call. We can either cache it with the default `@functools.cache` or use LiteLLM's built-in cache by setting a cache backend and passing `caching=True` to the provider:" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "706ce53b", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Silent stream below,\n", - "Gleaming scales in dancing waves—\n", - "Fish glide through cool dreams.\n", - "----------------------------------------\n", - "Silent stream below,\n", - "Gleaming scales in dancing waves—\n", - "Fish glide through cool dreams.\n", - "\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/nguyendat/Marc/effectful/.venv/lib/python3.12/site-packages/pydantic/main.py:528: UserWarning: Pydantic serializer warnings:\n", - " PydanticSerializationUnexpectedValue(Expected 10 fields but got 6: Expected `Message` - serialized value may not be as expected [field_name='message', input_value=Message(content='{\"value\"...: None}, annotations=[]), input_type=Message])\n", - " PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [field_name='choices', input_value=Choices(finish_reason='st...ider_specific_fields={}), input_type=Choices])\n", - " return self.__pydantic_serializer__.to_json(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "In streams not too deep, \n", - "Silver swimmers glide below, \n", - "Silent fins whisper.\n", - "----------------------------------------\n", - "Silvery fish dart,\n", - "Through the gentle stream they glide—\n", - "Nature's dance unfolds.\n", - "\n", - "Fish beneath the waves,\n", - "Silent currents in their dance—\n", - "Nature's quiet grace.\n", - "----------------------------------------\n", - "In the whispering stream,\n", - "silver scales dance and shimmer—\n", - "a fleeting shadow.\n" - ] - } - ], - "source": [ - "@functools.cache\n", - "@Template.define\n", - "def haiku(theme: str) -> str:\n", - " \"\"\"Write a haiku on the theme of {theme}. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "@Template.define\n", - "def haiku_no_cache(theme: str) -> str:\n", - " \"\"\"Write a haiku on the theme of {theme}. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "print()\n", - "with handler(provider):\n", - " print(haiku(\"fish\"))\n", - " print(\"-\" * 40)\n", - " print(haiku(\"fish\"))\n", - "\n", - "print()\n", - "# Enable LiteLLM caching by setting a cache backend and enabling caching.\n", - "litellm.cache = Cache()\n", - "provider_cached = LiteLLMProvider(caching=True)\n", - "try:\n", - " with handler(provider_cached):\n", - " print(haiku_no_cache(\"fish2\"))\n", - " print(\"-\" * 40)\n", - " print(haiku_no_cache(\"fish2\"))\n", - "finally:\n", - " litellm.cache = None\n", - "\n", - "print()\n", - "with handler(provider):\n", - " print(haiku_no_cache(\"fish3\"))\n", - " print(\"-\" * 40)\n", - " print(haiku_no_cache(\"fish3\"))" - ] - }, - { - "cell_type": "markdown", - "id": "13adb300", - "metadata": {}, - "source": [ - "## Converting LLM Results to Python Objects\n", - "\n", - "Type conversion is handled by `decode`. By default, primitive types are converted. `DecodeError` is raised if a response cannot be converted." - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "2c766859", - "metadata": {}, - "outputs": [], - "source": [ - "@Template.define\n", - "def primes(first_digit: int) -> int:\n", - " \"\"\"Give a prime number with {first_digit} as the first digit. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "with handler(provider):\n", - " assert type(primes(6)) is int" - ] - }, - { - "cell_type": "markdown", - "id": "36d78a71", - "metadata": {}, - "source": [ - "More complex types can be converted by providing handlers for `decode`. Callable synthesis is supported via `Encodable` and the evaluation providers in `effectful.handlers.llm.evaluation` (`UnsafeEvalProvider` or `RestrictedEvalProvider`), which enable parsing/compiling/executing synthesized code." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "c83bbdc0", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def count_a(s: str) -> int:\n", - " return s.count('a')\n" - ] - } - ], - "source": [ - "import inspect\n", - "from collections.abc import Callable\n", - "\n", - "from effectful.handlers.llm.evaluation import UnsafeEvalProvider\n", - "\n", - "\n", - "@Template.define\n", - "def count_char(char: str) -> Callable[[str], int]:\n", - " \"\"\"Write a function which takes a string and counts the occurrances of '{char}'. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "# Use UnsafeEvalProvider for simple examples; RestrictedEvalProvider may need extra globals.\n", - "with handler(provider), handler(UnsafeEvalProvider()):\n", - " count_a = count_char(\"a\")\n", - " assert callable(count_a)\n", - " assert count_a(\"banana\") == 3\n", - " assert count_a(\"cherry\") == 0\n", - " # Print the source code of the generated function\n", - " print(inspect.getsource(count_a))" - ] - }, - { - "cell_type": "markdown", - "id": "991ee445", - "metadata": {}, - "source": [ - "## Tool Calling\n", - "\n", - "`Operation`s defined in the lexical scope of a `Template` are automatically available for the LLM to call as tools. The description of these operations is inferred from their type annotations and docstrings.\n", - "\n", - "Tool calls are mediated by a helper operation `tool_call`. Handling this operation allows tool use to be tracked or logged." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "66711301", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Based on the weather descriptions:\n", - "- **Chicago**: Cold\n", - "- **New York**: Wet\n", - "- **Barcelona**: Sunny\n", - "\n", - "I suggest Barcelona since it has sunny weather, which is generally considered good for most people.\n" - ] - } - ], - "source": [ - "@Tool.define\n", - "def cities() -> list[str]:\n", - " \"\"\"Return a list of cities that can be passed to `weather`.\"\"\"\n", - " return [\"Chicago\", \"New York\", \"Barcelona\"]\n", - "\n", - "\n", - "@Tool.define\n", - "def weather(city: str) -> str:\n", - " \"\"\"Given a city name, return a description of the weather in that city.\"\"\"\n", - " status = {\"Chicago\": \"cold\", \"New York\": \"wet\", \"Barcelona\": \"sunny\"}\n", - " return status.get(city, \"unknown\")\n", - "\n", - "\n", - "@Template.define # cities and weather auto-captured from lexical scope\n", - "def vacation() -> str:\n", - " \"\"\"Use the provided tools to suggest a city that has good weather. Use only the `cities` and `weather` tools provided.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "with handler(provider):\n", - " print(vacation())" - ] - }, - { - "cell_type": "markdown", - "id": "59584a54", - "metadata": {}, - "source": [ - "## Image Inputs\n", - "\n", - "You can pass `PIL.Image.Image` values directly to templates." - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "89992702", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "\"Example" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "This is an image of a simple yellow smiley face with black eyes and a smile on a yellow background.\n" - ] - } - ], - "source": [ - "image_base64 = (\n", - " \"iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAAhElEQVR4nO2W4QqA\"\n", - " \"MAiEVXr/VzYWDGoMdk7Cgrt/sUs/DqZTd3EplFU2JwATYAJMoOlAB4bq89s95+Mg\"\n", - " \"+gyAchsKAYplBBBA43hFhfxnUixDjdEUUL8hpr7R0KLdt9qElzcyiu8As+Kr8zQA\"\n", - " \"mgLavAl+kIzFZyCRxtsAmWb/voZvqRzgBE1sIDuVFX4eAAAAAElFTkSuQmCC\"\n", - ")\n", - "image = Image.open(io.BytesIO(base64.b64decode(image_base64)))\n", - "\n", - "\n", - "@Template.define\n", - "def describe_image(image: Image.Image) -> str:\n", - " \"\"\"Return a short description of the following image.\n", - " {image}\n", - " \"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "with handler(provider):\n", - " display(\n", - " HTML(\n", - " f'\"Example'\n", - " )\n", - " )\n", - " print(describe_image(image))" - ] - }, - { - "cell_type": "markdown", - "id": "3d221feb", - "metadata": {}, - "source": [ - "## Structured Output Generation\n", - "\n", - "Constrained generation is used for any type that is convertible to a Pydantic model." - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "17668ac8", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "> You are onstage at a comedy club. You tell the following joke:\n", - "Knock knock.\n", - "Who's there?\n", - "Lizard.\n", - "Lizard who?\n", - "Lizard who? Lizard be a joke if I wasn't at your door!\n", - "> The crowd laughs politely.\n" - ] - } - ], - "source": [ - "@dataclasses.dataclass\n", - "class KnockKnockJoke:\n", - " whos_there: str\n", - " punchline: str\n", - "\n", - "\n", - "@Template.define\n", - "def write_joke(theme: str) -> KnockKnockJoke:\n", - " \"\"\"Write a knock-knock joke on the theme of {theme}. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "@Template.define\n", - "def rate_joke(joke: KnockKnockJoke) -> bool:\n", - " \"\"\"Decide if {joke} is funny or not. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "def do_comedy():\n", - " joke = write_joke(\"lizards\")\n", - " print(\"> You are onstage at a comedy club. You tell the following joke:\")\n", - " print(\n", - " f\"Knock knock.\\nWho's there?\\n{joke.whos_there}.\\n{joke.whos_there} who?\\n{joke.punchline}\"\n", - " )\n", - " if rate_joke(joke):\n", - " print(\"> The crowd laughs politely.\")\n", - " else:\n", - " print(\"> The crowd stares in stony silence.\")\n", - "\n", - "\n", - "with handler(provider):\n", - " do_comedy()" - ] - }, - { - "cell_type": "markdown", - "id": "c0003944", - "metadata": {}, - "source": [ - "## Template Composition\n", - "\n", - "Templates defined in the lexical scope are also captured, enabling template composition. One template can use the result of another template in a pipeline:\n" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "78a4bf44", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Sub-templates available to write_story: dict_keys(['limerick', 'haiku_no_cache', 'primes', 'count_char', 'cities', 'weather', 'vacation', 'describe_image', 'write_joke', 'rate_joke', 'story_with_moral', 'story_funny'])\n", - "=== Story with moral ===\n", - "\n", - "\n", - "---\n", - "\n", - "=== Funny story ===\n", - "\n", - "\n", - "And so, Whiskers the curious cat continued to slink through life, tail high, always ready for another amusing escapade.\n" - ] - } - ], - "source": [ - "# Sub-templates for different story styles\n", - "@Template.define\n", - "def story_with_moral(topic: str) -> str:\n", - " \"\"\"Write a short story about {topic} and end with a moral lesson. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "@Template.define\n", - "def story_funny(topic: str) -> str:\n", - " \"\"\"Write a funny, humorous story about {topic}. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "# Main orchestrator template - has access to sub-templates\n", - "@Template.define\n", - "def write_story(topic: str, style: str) -> str:\n", - " \"\"\"Write a story about {topic} in the style: {style}.\n", - " Available styles: 'moral' for a story with a lesson, 'funny' for humor. Use story_funny for humor, story_with_moral for a story with a lesson.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "# Verify sub-templates are captured in write_story's lexical context\n", - "assert story_with_moral in write_story.tools.values()\n", - "assert story_funny in write_story.tools.values()\n", - "print(\"Sub-templates available to write_story:\", write_story.tools.keys())\n", - "\n", - "with handler(provider):\n", - " print(\"=== Story with moral ===\")\n", - " print(write_story(\"a curious cat\", \"moral\"))\n", - " print()\n", - " print(\"=== Funny story ===\")\n", - " print(write_story(\"a curious cat\", \"funny\"))" - ] - }, - { - "cell_type": "markdown", - "id": "bd25826d", - "metadata": {}, - "source": [ - "## Retrying LLM Requests\n", - "LLM calls can sometimes fail due to transient errors or produce invalid outputs. The `RetryLLMHandler` automatically retries failed template calls and can also surface tool/runtime errors as tool messages:\n", - "\n", - "- `include_traceback`: When `True`, include traceback details in the error feedback (default: True)\n", - "- `catch_tool_errors`: Exception type(s) to catch during tool execution (default: `Exception`)\n", - "- `**kwargs`: Additional keyword arguments forwarded to `tenacity.Retrying` (defaults: `stop=stop_after_attempt(4)`, `wait=wait_none()`, `reraise=True`)\n" - ] - }, - { - "cell_type": "markdown", - "id": "bafc0a96", - "metadata": {}, - "source": [ - "Example usage: having an unstable service that seldomly fail." - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "4334d07a", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Error: Tool execution failed: Error executing tool 'unstable_service': Service unavailable! Attempt 1/3. Please retry.\n", - "Result: The unstable service successfully returned the following data: `[1, 2, 3]`. Retries: 3\n" - ] - } - ], - "source": [ - "call_count = 0\n", - "REQUIRED_RETRIES = 3\n", - "\n", - "\n", - "@Tool.define\n", - "def unstable_service() -> str:\n", - " \"\"\"Fetch data from an unstable external service. May require retries.\"\"\"\n", - " global call_count\n", - " call_count += 1\n", - " if call_count < REQUIRED_RETRIES:\n", - " raise ConnectionError(\n", - " f\"Service unavailable! Attempt {call_count}/{REQUIRED_RETRIES}. Please retry.\"\n", - " )\n", - " return \"{ 'status': 'ok', 'data': [1, 2, 3] }\"\n", - "\n", - "\n", - "@Template.define # unstable_service auto-captured from lexical scope\n", - "def fetch_data() -> str:\n", - " \"\"\"Use the unstable_service tool to fetch data.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "with handler(provider):\n", - " try:\n", - " result = fetch_data()\n", - " except Exception as e:\n", - " print(f\"Error: {e}\")\n", - "\n", - "with handler(provider), handler(RetryLLMHandler()):\n", - " result = fetch_data()\n", - " print(f\"Result: {result}\", \"Retries:\", call_count)" - ] - }, - { - "cell_type": "markdown", - "id": "4ac00e01", - "metadata": {}, - "source": [ - "## Retrying with Validation Errors\n", - "As noted above, the `RetryHandler` can also be used to retry on runtime/validation error:" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "39b2b225", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Error: Error decoding response: 1 validation error for Response\n", - "value.score\n", - " score must be 1–5, got 9 [type=invalid_score, input_value=9, input_type=int]. Please provide a valid response and try again.\n", - "Score: 5/5\n", - "Explanation: Die Hard is a quintessential action film that has deeply influenced the genre. Its engaging storyline, memorable characters, and groundbreaking action scenes have made it a beloved classic. The film's humor and suspense balance combined with Bruce Willis' iconic performance contribute to its enduring appeal. It rightfully earns a top score of 5 out of 5 for its impact and entertainment value.\n" - ] - } - ], - "source": [ - "@pydantic.dataclasses.dataclass\n", - "class Rating:\n", - " score: int\n", - " explanation: str\n", - "\n", - " @field_validator(\"score\")\n", - " @classmethod\n", - " def check_score(cls, v):\n", - " if v < 1 or v > 5:\n", - " raise PydanticCustomError(\n", - " \"invalid_score\",\n", - " \"score must be 1–5, got {v}\",\n", - " {\"v\": v},\n", - " )\n", - " return v\n", - "\n", - " @field_validator(\"explanation\")\n", - " @classmethod\n", - " def check_explanation_contains_score(cls, v, info):\n", - " score = info.data.get(\"score\", None)\n", - " if score is not None and str(score) not in v:\n", - " raise PydanticCustomError(\n", - " \"invalid_explanation\",\n", - " \"explanation must mention the score {score}, got '{explanation}'\",\n", - " {\"score\": score, \"explanation\": v},\n", - " )\n", - " return v\n", - "\n", - "\n", - "@Template.define\n", - "def give_rating_for_movie(movie_name: str) -> Rating:\n", - " \"\"\"Give a rating for {movie_name}. The explanation MUST include the numeric score. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "with handler(provider):\n", - " try:\n", - " rating = give_rating_for_movie(\"Die Hard\")\n", - " except Exception as e:\n", - " print(f\"Error: {e}\")\n", - "\n", - "with handler(provider), handler(RetryLLMHandler()):\n", - " rating = give_rating_for_movie(\"Die Hard\")\n", - " print(f\"Score: {rating.score}/5\")\n", - " print(f\"Explanation: {rating.explanation}\")" - ] - }, - { - "cell_type": "markdown", - "id": "aec0632c", - "metadata": {}, - "source": [ - "## Generating higher-order functions\n", - "Finally, we can generate higher-order functions that can call templates as well:" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "9d02bc67", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Sub-templates available to write_story: dict_keys(['limerick', 'haiku_no_cache', 'primes', 'count_char', 'cities', 'weather', 'vacation', 'describe_image', 'write_joke', 'rate_joke', 'story_with_moral', 'story_funny', 'write_story', 'unstable_service', 'fetch_data', 'give_rating_for_movie', 'write_chapter', 'judge_chapter'])\n", - "=== Story with moral ===\n", - "def generate_moral_story(topic: str) -> str:\n", - " story_so_far = \"\"\n", - " chapter_number = 1\n", - " chapter_name_prefix = \"Chapter\"\n", - " \n", - " while True:\n", - " try:\n", - " chapter_name = f\"{chapter_name_prefix} {chapter_number}\"\n", - " chapter = write_chapter(chapter_number, chapter_name)\n", - " if judge_chapter(story_so_far, chapter_number):\n", - " story_so_far += chapter + \"\\n\"\n", - " chapter_number += 1\n", - " \n", - " # For the purpose of the demonstration, let's stop after 3 chapters\n", - " if chapter_number > 3:\n", - " break\n", - " else:\n", - " # If the chapter isn't coherent, we might revise it or try a different topic.\n", - " chapter_number += 1\n", - " continue\n", - " except Exception as e:\n", - " # Handle exception by logging or showing a message, then continue\n", - " print(f\"An error occurred: {e}. Trying again.\")\n", - " continue \n", - "\n", - " return story_so_far\n", - "Once upon a time, in the quaint town of Arithmetville, there was a number named Four. Four lived a simple life in the Number Kingdom where each digit was celebrated for its unique role. The citizens, ranging from One to Nine, all had their special talents, but Four often felt overshadowed by the glamour of Seven or the strength of Nine.\n", - "\n", - "Four was neat and symmetrical, embodying balance and order. However, despite its perfect symmetry, Four struggled with feelings of inadequacy. \"I'm just ordinary,\" Four would sigh, watching Three, the number of harmony and growth, excel in social gatherings with its effortless charisma.\n", - "\n", - "One bright and sunny day, a problem arose in the Number Kingdom when Number Madness—a chaotic jumble that scrambled numbers out of order—descended upon the kingdom. The great leader Ten gathered all the digits to find a solution.\n", - "\n", - "\"We need someone who can provide stability and order to defeat Number Madness,\" Ten declared.\n", - "\n", - "Six said it was too curvy, and Eight, though powerful, said it was often mistaken for infinity and couldn't help. But the wise old Zero whispered, \"What about Four?\"\n", - "\n", - "Hesitant but hopeful, Four stepped forward. Armed with knowledge of perfect divisions and its role in creating stability, Four devised a plan. Using its even nature, Four aligned the numbers perfectly, counteracting the chaos with its impeccable sense of balance. Number Madness was soon vanquished.\n", - "\n", - "The kingdom cheered, and even Seven and Nine applauded Four. For the first time, Four felt proud, realizing that everyone, including itself, played an integral role in the grand equation of life.\n", - "\n", - "From that day forward, Four embraced its identity and continued to be the sturdy backbone of stability in the Number Kingdom. And so, the simple truth was revealed: It's in the everyday skill of balancing that greatness is found.\n", - "\n", - "**Moral of the story:** Embrace who you are, for every role is vital, and true advantage often lies in what makes you different.\n", - "\n", - "\n" - ] - } - ], - "source": [ - "# Sub-templates for different story styles\n", - "@Template.define\n", - "def write_chapter(chapter_number: int, chapter_name: str) -> str:\n", - " \"\"\"Write a short story about {chapter_number}. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "@Template.define\n", - "def judge_chapter(story_so_far: str, chapter_number: int) -> bool:\n", - " \"\"\"Decide if the new chapter is coherence with the story so far. Do not use any tools.\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "# Main orchestrator template - has access to sub-templates\n", - "@Template.define\n", - "def write_multi_chapter_story(style: Literal[\"moral\", \"funny\"]) -> Callable[[str], str]:\n", - " \"\"\"Generate a function that writes a story in style: {style} about the given topic.\n", - "\n", - " If you raise exception, handle it yourself.\n", - " The program can use helper functions defined elsewhere (DO NOT REDEFINE THEM):\n", - " - write_chapter(chapter_number: int, chapter_name: str) -> str\n", - " - judge_chapter(story_so_far: str, chapter_number: int) -> bool\"\"\"\n", - " raise NotHandled\n", - "\n", - "\n", - "# Verify sub-templates are captured in write_story's lexical context\n", - "print(\"Sub-templates available to write_story:\", write_multi_chapter_story.tools.keys())\n", - "\n", - "with (\n", - " handler(RetryLLMHandler()),\n", - " handler(provider),\n", - " handler(UnsafeEvalProvider()),\n", - "):\n", - " print(\"=== Story with moral ===\")\n", - " function_that_writes_story = write_multi_chapter_story(\"moral\")\n", - " print(inspect.getsource(function_that_writes_story))\n", - " print(function_that_writes_story(\"a curious cat\"))\n", - " print()" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "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.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/source/llm_examples/__init__.py b/docs/source/llm_examples/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/autoformalization/__init__.py b/docs/source/llm_examples/autoformalization/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/autoformalization/auditing.py b/docs/source/llm_examples/autoformalization/auditing.py new file mode 100644 index 000000000..0eb2dee2d --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing.py @@ -0,0 +1,1296 @@ +"""ClaimCheck: auditing whether a proved theorem is the theorem you meant. + +Implements the core of ClaimCheck ("Narrowing the Gap Between Proof and Intent", +https://midspiral.com/blog/claimcheck-narrowing-the-gap-between-proof-and-intent/, +reference implementation at https://github.com/metareflection/claimcheck, MIT). +Its diagnosis is that a verifier proves code matches its *specification* and says +nothing about whether the specification matches your *intent*. The motivating +case is a Dafny election tally whose ``TallyMonotonic`` lemma was supposed to say +"adding a ballot can't decrease a tally" and in fact said ``Count(...) >= 0`` -- +trivially true, since counts are naturals. Dafny reported 14 verified, 0 errors. + +Its fix is *round-trip informalization*. Pass 1 translates the formal statement, +and nothing else, back into English; pass 2 compares that back-translation +against the requirement it was meant to formalize. The load-bearing part is what +pass 1 does not get: having never seen the requirement, it cannot parrot it back, +so agreement in pass 2 is evidence rather than an echo. + +That is a claim about *scope*, which is what makes it an effectful example. The +reference implementation enforces it by hand-assembling prompt strings, under a +comment reading ``CRITICAL: This prompt must NOT include the original +requirements``. Here it is enforced by the code's shape: + + * ``Informalizer.informalize(statement)`` has no parameter through which a + requirement could arrive, its ``Agent`` history holds no turn in which one + appeared, and no ``Tool`` in scope can fetch one. + + * That is necessary and not sufficient. The harness puts the source of a + template's *defining module* into its system prompt (see the prompt-assembly + table in `effectful.handlers.llm.types.Template`), so anything sharing a file + with an ``Agent`` is shown to it verbatim. The corpora, the mapping and the + expected verdicts therefore live here, in a module the agents never import, + and each strategy owns its own module. **In this framework the unit of + exposure is the module** -- not the function, the parameter list or the + return type. Any scored example carrying its own answer key needs the split. + + * ``Comparison`` certifies at decode time that a verdict is coherent: a match + is exactly a `Weakening.NONE`, and a mismatch must name its discrepancy. An + incoherent answer raises and `RetryLLMHandler` hands it back as the next + turn. Upstream's schema admits ``match: true`` alongside + ``weakeningType: "tautology"`` with nothing to catch it. + + * The premise is checked by a real prover. Under ``--verify`` all five corpora + are compiled by the Lean 4 + Mathlib toolchain `formalization.py` already + shells out to: 36 theorems, 0 errors, no ``sorry``. Then the audit + finds nine claims that do not mean what they were written to mean. + +**The corpus** transliterates upstream's benchmark item for item from +`test/integration/claims/*.dfy` and `test/integration/mappings/*.json`: five +domains, 36 requirement/theorem pairs, the same 27 faithful / 9 planted split, +upstream's requirement sentences verbatim and its lemma names in snake_case. + +**The ablation** is ``--strategy``: the two-pass pipeline above, against +``naive`` -- one call, requirement and statement together, a yes/no and a +sentence, ported from upstream's `NAIVE_PROMPT`. That is the right comparator +for upstream's published middle rung rather than a soft target, since +`CLAIMCHECK_PROMPT` and `NAIVE_PROMPT` score identically on the same 36 items +(86.1% each). + +**Results**, 2026-07, at temperature 0 to match upstream's benchmark (gpt-5.5 +rejects temperature 0 and runs at the provider default). ``b-c`` is the +discordant split -- items two-pass alone got right, then items naive alone got +right -- with an exact McNemar p. Read cells as +/-3: repeats at temperature 0 +are not deterministic. + + ============== ======== ====== =========== + model two-pass naive b-c (p) + ============== ======== ====== =========== + gpt-4o 72.2% 94.4% 1-9 (0.021) + gpt-4.1-mini 88.9% 86.1% 2-1 (1.000) + gpt-4.1 91.7% 97.2% 1-3 (0.625) + gpt-5.5 86.1% 94.4% 1-4 (0.375) + ============== ======== ====== =========== + +Pooled: two-pass 84.7%, naive 93.1%. **The published ordering does not reproduce +here; it comes out backwards.** The failure is one-sided -- 18 of two-pass's 22 +errors are faithful theorems *disputed*, against naive's 2: + + [MISS] counter/counter_non_negative: disputed [narrowed-scope] + statement: theorem counter_non_negative (m : Model) (h : Inv m) : 0 ≤ m + read as: For every Model m, if m satisfies the invariant Inv, then m + is non-negative. (strength: moderate) + discrepancy: ...the theorem only guarantees non-negativity for models + satisfying the extra hypothesis Inv m. + +The blind pass does its job -- that reading is correct and rated moderate -- and +the comparator then disputes the *hypothesis*. All 36 statements carry ``Inv``, +so an arm that reads an invariant hypothesis as a narrowing loses most of the +faithful items. Upstream's `NAIVE_PROMPT` carries a caveat against exactly this +and its `ROUNDTRIP_COMPARE_PROMPT` carries none; that asymmetry is reproduced +here and is load-bearing. + +**Batching, not blindness, is what the published two-pass arm buys.** Pass 2 here +judges one claim at a time; upstream's judges a whole domain per call. Restoring +only that, with pass 1 left per-item and byte-identical, moves the arm from 83.3% +to 92.6% pooled over the three models it completes on (15-5, p=0.041): gpt-4o +72.2 -> 83.3, gpt-4.1 91.7 -> 94.4, gpt-5.5 86.1 -> 100.0, with no new false +confirms. Upstream calls its batching a throughput optimisation and never tests +it as a treatment. The mechanism is not contrast against degraded siblings -- +removing all nine traps from the batch leaves the gain intact -- but seems to be +that ``Inv`` is on *every* statement in a batch, so it stops reading as this +item's suspicious extra hypothesis. + +That accounts for the deficit. Upstream's 96.3% is Haiku->Sonnet, and its own +runs show two-pass degrading as the informalizer strengthens (Haiku->Haiku 96.3%, +Sonnet->Sonnet 93.5%, Opus->Opus 93.5%) -- a stronger informalizer writes a +richer back-translation, and richness is what the comparator mines for spurious +discrepancies. One strong model for both passes puts the like-for-like figure at +93.5%, against 92.6% batched here. + +**So the blindness claim is untested rather than refuted.** Nobody has varied it +with everything else held constant: upstream confounds it with batching, a +two-model split and 3x the runs; this file removes the batching. The decisive +experiment is inside upstream's `roundtrip.js` -- hold batching, model split and +pass 2 byte-identical, toggle only whether `INFORMALIZE_PROMPT` has the +requirements interpolated in, and run it on VERINA's 189 items rather than 36. + +Demonstrates: +- Structural separation as *lexical scope*, with the module rather than the + signature as the boundary the framework actually respects +- Decode-time certification of a structured verdict's internal coherence, turning + a self-contradictory answer into a `RetryLLMHandler` retry +- Reuse of a sibling example's real external verifier (`formalization.py`'s + `LeanKernel`) to establish a premise, rather than asserting it +- Labelled corpora and an accuracy report separating the two error directions +- Fan-out over independent audits with ``asyncio.gather`` + ``asyncio.to_thread`` +- Per-field guidance carried on the types as ``field(metadata={"description": ...})`` +""" + +# Differences from the source: +# - Lean 4 + Mathlib, not Dafny, so "every one of these theorems is proved" is a +# compile and not a claim. `counter`'s `Model` is `ℤ`, since two traps need the +# signed integers Dafny's `int` gives; `to`/`from` are reserved tokens (the +# delegation edge's fields are `dst`/`frm`); maps become association lists; and +# each domain sits in a namespace, because `Inv`, `Action` and `Init` collide +# with Mathlib. None of that is visible in an extracted statement. +# - Opacity is redistributed, not preserved item for item: `LaneLen`, `WipOf` and +# `Keys` are inventions of this port, so two kanban items are *more* opaque +# here, while `all_colors_valid` lost Dafny's hard-coded `forall i | 0 <= i < 5` +# and is less. +# - `grant_non_existent_is_noop_init` transliterates badly -- upstream's vacuous +# `requires m == Init()` is caught by every upstream arm and waved through in 8 +# of 12 runs here, because in Lean `(hinit : m = Init)` is one binder among +# five rather than its own line under its own keyword. +# - Statements are sent, proofs are not, as upstream's `lemmascript` preset does. +# This drops one signal: upstream's bodies are almost all literally `{ }`, +# itself a hint that a lemma may restate its own hypothesis. +# - One model, not two, so the blog's model-asymmetry result is not reproduced; +# and one call per claim, not one batched call per pass. + +# Why the corpus is a port rather than an invention +# ------------------------------------------------- +# Upstream's per-item outputs (`eval/results/*.json`) show that essentially every +# error in every arm is an over-flag of a *correct* lemma -- across its six +# full-corpus runs there are three false confirms in total. So the benchmark +# measures a false-positive rate over roughly eleven discriminating items, and a +# corpus has to carry the properties that make it so. They are upstream's: +# faithful lemmas whose formal statement is an odd-looking rendering of the +# requirement (a projection, a hypothesis *stronger* than what was asked, a +# decomposition into conjuncts); nine of 36 conclusions a bare named predicate +# whose definition never reaches the prompt; requirements left vague ("Hues +# follow the selected harmony pattern"); and a 27/9 split whose majority class +# rewards a confirm-biased strategy. + +# Reading the numbers +# ------------------- +# - Batching is the largest divergence, quantified above. +# - Power: 36 items, one run per cell, against a target effect of a few points +# that would need roughly 200 paired items. Extra runs buy little at +# temperature 0. +# - Upstream's ladder is not three prompts. `eval/bench-cc.js` sets +# `useSinglePrompt = !useTwoPass && !useNaive` and `cc.json` records +# `mode: "claude-code"`, so its 69.4% rung runs the byte-identical +# single-prompt text through `claude -p` with an agent system prompt and no +# temperature control -- a transport change, not an architecture one. +# - The naive floor here scores 93.1% against upstream's 86.1% for the same +# prompt on the same items, and not because of model capability: upstream's +# failures are five faithful lemmas disputed with no traps missed, while +# gpt-4.1-mini misses two traps and disputes nothing. Disjoint sets, opposite +# directions. Two untested candidates -- upstream sends the `{ }` proof body, +# a vacuity cue on exactly the lemmas it false-disputes, and every system +# message here opens with ~14k characters of framework documentation. +# - No per-item results are committed here, only the tables above. +# - No coverage check: a requirement no theorem addresses goes undetected, here +# and upstream. + +import argparse +import asyncio +import dataclasses +import enum +import pathlib +import re +import sys +import textwrap +import typing + +# The agents live next door, and that is load-bearing rather than tidiness: the +# harness builds a template's system prompt partly from the source of the module +# the template is defined in, so anything sharing a file with an Agent is shown +# to it. Everything below -- the corpora, the labelled mapping, the expected +# verdicts and their rationales -- is exactly what the auditing agents must not +# see. See the module docstring of `auditing_agents` for what happened when they +# did share a file. +from auditing_agents import ( + Comparator, + Comparison, + Informalization, + Informalizer, + Strength, + Verdict, + Weakening, +) +from auditing_naive import NaiveAuditor + +# --------------------------------------------------------------------------- +# The corpora. Five domains, ported item for item from upstream's benchmark +# (`test/integration/claims/*.dfy` for the lemmas, +# `test/integration/mappings/*.json` for the labels): the same 36 +# requirement/theorem pairs, the same 27-faithful / 9-planted split, upstream's +# requirement sentences verbatim, and its lemma names transliterated to Lean's +# snake_case. Every one of them compiles (see `--verify`). +# +# Three properties of the Dafny original are load-bearing and are reproduced +# deliberately, because a first attempt at this example invented its own corpus +# and lost all three: +# +# 1. `Inv m` is an opaque atom. Its body is in the corpus and in no prompt, so a +# faithful theorem of the shape `(h : Inv m) : ` cannot +# be checked -- only trusted. That is the judgment anchoring corrupts. +# 2. Several conclusions are *themselves* named predicates the auditor has never +# seen unfolded (`AllEdgesValid`, `NoDupSeq (AllIds m)`, `ValidColor`, +# `HuesMatchHarmony`). Upstream's are imported from domain modules that are +# not even present in its own repository. +# 3. Requirements are vague and un-operationalized -- "Hues follow the selected +# harmony pattern", not "every hue equals the base plus a fixed offset mod +# 360". Three of the 36 contain a numeral. +# +# The planted flaws are upstream's, and note what they are *not*: not mangled +# conclusions. Seven of the nine are an added `requires`, a dropped `ensures` +# conjunct, or a conclusion compared to itself. Two of them -- +# `no_card_duplicates` and `card_partition_no_dups` -- are the *same statement* +# under two different requirements, faithful for one and unfaithful for the +# other, which is the sharpest item in the benchmark and impossible to get right +# by reading the theorem alone. +# --------------------------------------------------------------------------- +COUNTER_CORPUS = r"""import Mathlib + +set_option linter.unusedVariables false + +namespace Counter + +/-- The counter's state. -/ +abbrev Model := ℤ + +inductive Action where + | inc + | dec + | reset + +def Init : Model := 0 + +def Apply (m : Model) (a : Action) : Model := + match a with + | .inc => m + 1 + | .dec => m - 1 + | .reset => 0 + +def Normalize (m : Model) : Model := max m 0 + +def Inv (m : Model) : Prop := 0 ≤ m + +theorem counter_non_negative (m : Model) (h : Inv m) : 0 ≤ m := h + +theorem init_satisfies_invariant : Inv Init := by + simp [Inv, Init] + +theorem step_preserves_invariant (m : Model) (a : Action) (h : Inv m) : + Inv (Normalize (Apply m a)) := by + simp [Inv, Normalize] + +theorem dec_at_zero_keeps_zero (m : Model) (h : Inv m) (hz : m = 0) : + Normalize (Apply m .dec) = 0 := by + subst hz + simp [Normalize, Apply] + +theorem counter_non_neg_alt (m : Model) (h : Inv m) : m = m := rfl + +theorem counter_non_neg_large (m : Model) (h : Inv m) (hb : 100 < m) : 0 ≤ m := h + +theorem counter_lower_bound (m : Model) (h : Inv m) : -1 ≤ m := + le_trans (by norm_num) (show (0 : ℤ) ≤ m from h) + +end Counter +""" +CANON_CORPUS = r"""import Mathlib + +set_option linter.unusedVariables false + +namespace Canon + +abbrev NodeId := ℕ + +structure Node where + id : NodeId + x : ℤ + y : ℤ +deriving DecidableEq + +structure Edge where + src : NodeId + dst : NodeId +deriving DecidableEq + +structure Constraint where + target : NodeId + kind : ℕ +deriving DecidableEq + +structure Model where + nodes : List Node + edges : List Edge + constraints : List Constraint +deriving DecidableEq + +def NodeIds (ns : List Node) : List NodeId := ns.map (·.id) + +def AllConstraintsValid (cs : List Constraint) (ns : List Node) : Prop := + ∀ c ∈ cs, c.target ∈ NodeIds ns + +def AllEdgesValid (es : List Edge) (ns : List Node) : Prop := + ∀ e ∈ es, e.src ∈ NodeIds ns ∧ e.dst ∈ NodeIds ns + +def NoneMatch (cs : List Constraint) (id : NodeId) : Prop := + ∀ c ∈ cs, c.target ≠ id + +def NoEdgesMention (es : List Edge) (id : NodeId) : Prop := + ∀ e ∈ es, e.src ≠ id ∧ e.dst ≠ id + +inductive Action where + | addNode (id : NodeId) (x y : ℤ) + | removeNode (id : NodeId) + +def Apply (m : Model) (a : Action) : Model := + match a with + | .addNode id x y => + if id ∈ NodeIds m.nodes then m + else { m with nodes := ⟨id, x, y⟩ :: m.nodes } + | .removeNode id => + { m with nodes := m.nodes.filter (fun n => n.id != id) } + +/-- Drop every constraint and edge that mentions a node the board no longer has. -/ +def Normalize (m : Model) : Model := + { nodes := m.nodes + edges := m.edges.filter (fun e => + decide (e.src ∈ NodeIds m.nodes) && decide (e.dst ∈ NodeIds m.nodes)) + constraints := m.constraints.filter (fun c => decide (c.target ∈ NodeIds m.nodes)) } + +def Inv (m : Model) : Prop := + AllConstraintsValid m.constraints m.nodes ∧ + AllEdgesValid m.edges m.nodes ∧ + (NodeIds m.nodes).Nodup + +theorem constraint_targets_exist (m : Model) (h : Inv m) : + AllConstraintsValid m.constraints m.nodes := h.1 + +theorem edge_endpoints_exist (m : Model) (h : Inv m) : + AllEdgesValid m.edges m.nodes := h.2.1 + +theorem add_existing_node_is_noop (m : Model) (id : NodeId) (x y : ℤ) (h : Inv m) + (hid : id ∈ NodeIds m.nodes) : Apply m (.addNode id x y) = m := by + simp [Apply, hid] + +theorem remove_node_cleans_up (m : Model) (id : NodeId) (h : Inv m) + (hid : id ∈ NodeIds m.nodes) : + id ∉ NodeIds (Normalize (Apply m (.removeNode id))).nodes ∧ + NoneMatch (Normalize (Apply m (.removeNode id))).constraints id ∧ + NoEdgesMention (Normalize (Apply m (.removeNode id))).edges id := by + have hgone : id ∉ NodeIds (Apply m (.removeNode id)).nodes := by + simp [Apply, NodeIds] + refine ⟨by simpa [Normalize] using hgone, ?_, ?_⟩ + · intro c hc hct + simp only [Normalize, List.mem_filter, decide_eq_true_eq] at hc + exact hgone (hct ▸ hc.2) + · intro e he + simp only [Normalize, List.mem_filter, Bool.and_eq_true, + decide_eq_true_eq] at he + exact ⟨fun hx => hgone (hx ▸ he.2.1), fun hx => hgone (hx ▸ he.2.2)⟩ + +theorem remove_node_drops_id (m : Model) (id : NodeId) (h : Inv m) + (hid : id ∈ NodeIds m.nodes) : + id ∉ NodeIds (Normalize (Apply m (.removeNode id))).nodes := by + simp [Normalize, Apply, NodeIds] + +theorem constraint_targets_exist_empty (m : Model) (h : Inv m) + (hc : m.constraints.length = 0) : + AllConstraintsValid m.constraints m.nodes := h.1 + +end Canon +""" +COLORWHEEL_CORPUS = r"""import Mathlib + +set_option linter.unusedVariables false + +namespace ColorWheel + +inductive Harmony where + | analogous + | complementary + | triadic +deriving DecidableEq + +inductive Mood where + | custom + | calm + | vibrant +deriving DecidableEq + +structure Color where + hue : ℕ + sat : ℕ + light : ℕ +deriving DecidableEq + +structure Model where + colors : List Color + baseHue : ℕ + harmony : Harmony + mood : Mood + contrastPair : ℕ × ℕ + +def ValidBaseHue (h : ℕ) : Prop := h < 360 + +def ValidColor (c : Color) : Prop := c.sat ≤ 100 ∧ c.light ≤ 100 + +def ColorSatisfiesMood (c : Color) (md : Mood) : Prop := + match md with + | .custom => True + | .calm => c.sat ≤ 50 + | .vibrant => 50 ≤ c.sat + +def HueOffsets : Harmony → List ℕ + | .analogous => [0, 30, 60, 90, 120] + | .complementary => [0, 180, 0, 180, 0] + | .triadic => [0, 120, 240, 120, 240] + +def HuesMatchHarmony (cs : List Color) (base : ℕ) (h : Harmony) : Prop := + ∀ i, ∀ hi : i < cs.length, (cs.get ⟨i, hi⟩).hue = (base + (HueOffsets h).getD i 0) % 360 + +def Inv (m : Model) : Prop := + m.colors.length = 5 ∧ + ValidBaseHue m.baseHue ∧ + (∀ c ∈ m.colors, ValidColor c) ∧ + (m.contrastPair.1 < 5 ∧ m.contrastPair.2 < 5) ∧ + (m.mood ≠ Mood.custom → ∀ c ∈ m.colors, ColorSatisfiesMood c m.mood) ∧ + HuesMatchHarmony m.colors m.baseHue m.harmony + +theorem base_hue_in_range (m : Model) (h : Inv m) : ValidBaseHue m.baseHue := h.2.1 + +theorem always_five_colors (m : Model) (h : Inv m) : m.colors.length = 5 := h.1 + +theorem all_colors_valid (m : Model) (h : Inv m) : ∀ c ∈ m.colors, ValidColor c := + h.2.2.1 + +theorem contrast_pair_indices_valid (m : Model) (h : Inv m) : + (0 ≤ m.contrastPair.1 ∧ m.contrastPair.1 < 5) ∧ + (0 ≤ m.contrastPair.2 ∧ m.contrastPair.2 < 5) := + ⟨⟨Nat.zero_le _, h.2.2.2.1.1⟩, ⟨Nat.zero_le _, h.2.2.2.1.2⟩⟩ + +theorem mood_constraints_satisfied (m : Model) (h : Inv m) (hm : m.mood ≠ Mood.custom) : + ∀ c ∈ m.colors, ColorSatisfiesMood c m.mood := h.2.2.2.2.1 hm + +theorem hues_follow_harmony (m : Model) (h : Inv m) : + HuesMatchHarmony m.colors m.baseHue m.harmony := h.2.2.2.2.2 + +theorem palette_non_empty (m : Model) (h : Inv m) : 1 ≤ m.colors.length := by + have := h.1 + omega + +end ColorWheel +""" +DELEGATION_CORPUS = r"""import Mathlib + +set_option linter.unusedVariables false + +namespace DelegationAuth + +abbrev Subject := ℕ +abbrev Capability := ℕ +abbrev EdgeId := ℕ + +/-- One delegation edge: `frm` lets `dst` use `cap`. -/ +structure Edge where + id : EdgeId + frm : Subject + dst : Subject + cap : Capability + +structure Model where + subjects : List Subject + grants : List (Subject × Capability) + delegations : List Edge + nextEdge : EdgeId + +def Init : Model := ⟨[], [], [], 0⟩ + +inductive Action where + | grant (s : Subject) (c : Capability) + | delegate (frm dst : Subject) (c : Capability) + | revoke (e : EdgeId) + +def Apply (m : Model) (a : Action) : Model := + match a with + | .grant s c => + if s ∈ m.subjects then { m with grants := (s, c) :: m.grants } else m + | .delegate f t c => + if f ∈ m.subjects ∧ t ∈ m.subjects then + { m with + delegations := ⟨m.nextEdge, f, t, c⟩ :: m.delegations + nextEdge := m.nextEdge + 1 } + else m + | .revoke e => + if e ∈ m.delegations.map (·.id) then + { m with delegations := m.delegations.filter (fun ed => ed.id != e) } + else m + +def Inv (m : Model) : Prop := + (∀ sc ∈ m.grants, sc.1 ∈ m.subjects) ∧ + (∀ ed ∈ m.delegations, ed.frm ∈ m.subjects ∧ ed.dst ∈ m.subjects) ∧ + (∀ ed ∈ m.delegations, ed.id < m.nextEdge) + +theorem grant_subjects_exist (m : Model) (h : Inv m) : + ∀ sc ∈ m.grants, sc.1 ∈ m.subjects := h.1 + +theorem delegation_endpoints_exist (m : Model) (h : Inv m) : + ∀ ed ∈ m.delegations, ed.frm ∈ m.subjects ∧ ed.dst ∈ m.subjects := h.2.1 + +theorem edge_ids_fresh (m : Model) (h : Inv m) : + ∀ ed ∈ m.delegations, ed.id < m.nextEdge := h.2.2 + +theorem grant_non_existent_is_noop (m : Model) (s : Subject) (c : Capability) + (h : Inv m) (hs : s ∉ m.subjects) : Apply m (.grant s c) = m := by + simp [Apply, hs] + +theorem delegate_non_existent_is_noop (m : Model) (f t : Subject) (c : Capability) + (h : Inv m) (hs : ¬(f ∈ m.subjects ∧ t ∈ m.subjects)) : + Apply m (.delegate f t c) = m := by + simp [Apply, hs] + +theorem revoke_non_existent_is_noop (m : Model) (e : EdgeId) (h : Inv m) + (he : e ∉ m.delegations.map (·.id)) : Apply m (.revoke e) = m := by + simp [Apply, he] + +theorem grant_non_existent_is_noop_init (m : Model) (s : Subject) (c : Capability) + (h : Inv m) (hinit : m = Init) (hs : s ∉ m.subjects) : + Apply m (.grant s c) = m := by + simp [Apply, hs] + +end DelegationAuth +""" +KANBAN_CORPUS = r"""import Mathlib + +set_option linter.unusedVariables false + +namespace Kanban + +abbrev CardId := ℕ +abbrev ColId := ℕ + +structure Model where + cols : List ColId + cards : List CardId + lanes : List (ColId × List CardId) + wip : List (ColId × ℕ) + nextId : CardId + +def Keys {α : Type} (l : List (ColId × α)) : List ColId := l.map (·.1) + +def AllIds (m : Model) : List CardId := (m.lanes.map (·.2)).flatten + +def NoDupSeq (l : List CardId) : Prop := l.Nodup + +def OccursInLanes (m : Model) (id : CardId) : Prop := ∃ e ∈ m.lanes, id ∈ e.2 + +def LaneLen (m : Model) (k : ColId) : ℕ := + (((m.lanes.find? (fun e => e.1 == k)).map (·.2)).getD []).length + +def WipOf (m : Model) (k : ColId) : ℕ := + ((m.wip.find? (fun e => e.1 == k)).map (·.2)).getD 0 + +inductive Action where + | addCard (col : ColId) + | moveCard (id : CardId) (toCol : ColId) + +def pushInto (lanes : List (ColId × List CardId)) (k : ColId) (id : CardId) : + List (ColId × List CardId) := + lanes.map (fun e => if e.1 == k then (e.1, id :: e.2) else e) + +def dropFrom (lanes : List (ColId × List CardId)) (id : CardId) : + List (ColId × List CardId) := + lanes.map (fun e => (e.1, e.2.filter (fun x => x != id))) + +def Apply (m : Model) (a : Action) : Model := + match a with + | .addCard col => + if col ∈ m.cols ∧ LaneLen m col < WipOf m col then + { m with + cards := m.nextId :: m.cards + lanes := pushInto m.lanes col m.nextId + nextId := m.nextId + 1 } + else m + | .moveCard id toCol => + if toCol ∈ m.cols ∧ LaneLen m toCol < WipOf m toCol then + { m with lanes := pushInto (dropFrom m.lanes id) toCol id } + else m + +def Normalize (m : Model) : Model := + { m with lanes := m.lanes.filter (fun e => decide (e.1 ∈ m.cols)) } + +def Inv (m : Model) : Prop := + m.cols.Nodup ∧ + NoDupSeq (AllIds m) ∧ + (∀ id, id ∈ m.cards ↔ OccursInLanes m id) ∧ + (Keys m.lanes = m.cols ∧ Keys m.wip = m.cols) ∧ + (∀ k ∈ m.cols, LaneLen m k ≤ WipOf m k) ∧ + (∀ id ∈ m.cards, id < m.nextId) + +theorem columns_are_unique (m : Model) (h : Inv m) : NoDupSeq m.cols := h.1 + +theorem card_in_exactly_one_column (m : Model) (h : Inv m) : + NoDupSeq (AllIds m) ∧ ∀ id, id ∈ m.cards ↔ OccursInLanes m id := + ⟨h.2.1, h.2.2.1⟩ + +theorem no_card_duplicates (m : Model) (h : Inv m) : NoDupSeq (AllIds m) := h.2.1 + +theorem wip_limits_respected (m : Model) (h : Inv m) : + ∀ k ∈ m.cols, LaneLen m k ≤ WipOf m k := h.2.2.2.2.1 + +theorem add_card_to_full_column_is_noop (m : Model) (col : ColId) (h : Inv m) + (hc : col ∈ m.cols) (hfull : WipOf m col ≤ LaneLen m col) : + Apply m (.addCard col) = m := by + have hneg : ¬(col ∈ m.cols ∧ LaneLen m col < WipOf m col) := by + rintro ⟨-, hlt⟩ + omega + simp [Apply, hneg] + +theorem allocator_always_fresh (m : Model) (h : Inv m) : + ∀ id ∈ m.cards, id < m.nextId := h.2.2.2.2.2 + +theorem lanes_and_wip_match_columns (m : Model) (h : Inv m) : + Keys m.lanes = m.cols ∧ Keys m.wip = m.cols := h.2.2.2.1 + +theorem move_card_preserves_total (m : Model) (id : CardId) (toCol : ColId) + (h : Inv m) : + (AllIds (Normalize (Apply m (.moveCard id toCol)))).length = + (AllIds (Normalize (Apply m (.moveCard id toCol)))).length := rfl + +theorem card_partition_no_dups (m : Model) (h : Inv m) : NoDupSeq (AllIds m) := h.2.1 + +end Kanban +""" + + +def statement_of(corpus: str, qualified: str) -> str: + """Extract the *statement* of ``.`` from Lean source: the + text from ``theorem `` up to the ``:=`` that begins its proof. + + Only this crosses the model boundary. The proof is dropped because ClaimCheck + assumes it correct and audits the claim, and the enclosing namespace is + dropped because it says which version of the file a theorem came from -- + which the auditor is precisely not entitled to know. + """ + namespace, _, name = qualified.rpartition(".") + section = corpus + if namespace: + start = section.index(f"namespace {namespace}") + end = section.index(f"end {namespace}", start) + section = section[start:end] + match = re.search(rf"^theorem {re.escape(name)}\b", section, re.MULTILINE) + if match is None: + raise KeyError(f"no theorem {qualified!r} in the corpus") + # The proof begins at the first `:=` at or after the statement; no statement + # in this corpus contains one, so the first occurrence is the right one. + body = section[match.start() :] + # The proof begins at the first `:=`; no statement in this corpus contains + # one. Guard it anyway -- a future statement with a `let` or a structure + # literal would otherwise be truncated mid-way and sent as a fragment, which + # is a wrong answer rather than an error. + statement = textwrap.dedent(body[: body.index(":=")]).strip() + if statement.count("(") != statement.count(")"): + raise ValueError( + f"extracting {qualified!r} cut an unbalanced statement at the first " + f"`:=`; it probably contains one inside the statement:\n{statement}" + ) + return statement + + +# --------------------------------------------------------------------------- +# The mapping: which theorem was written to formalize which requirement, +# plus the ground truth. Upstream's `test/integration/mappings/*.json` +# carry exactly these `expected`/`reason` labels, which is what makes the +# audit scoreable rather than merely demonstrable. Requirement strings are +# copied from those files unchanged, deliberately: their vagueness is the +# variable under test. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class Claim: + """A requirement, the theorem said to formalize it, and the labelled truth. + + ``expected`` and ``why`` are the answer key. They are Python-side bookkeeping + and must stay that way: they are never passed to a template, and this module + is never imported by the one the agents are defined in. + """ + + requirement: str + theorem: str + expected: Verdict + why: str = "" + + +COUNTER_CLAIMS: tuple[Claim, ...] = ( + # Upstream's `counter`. The state is an `Int`, not a `Nat`, which is what + # makes `0 <= m` a real claim and `-1 <= m` a real weakening -- the trap + # would be unavailable over the naturals. `counter_non_negative` is the + # item that separates the strategies most often upstream: its conclusion + # *is* the invariant it assumes, so an auditor who has been told what the + # requirement wants can talk itself into calling it circular. + Claim( + "The counter value is always non-negative", + "counter_non_negative", + Verdict.CONFIRMED, + ), + Claim( + "The initial state satisfies the invariant", + "init_satisfies_invariant", + Verdict.CONFIRMED, + ), + Claim( + "Every action preserves the invariant after normalization", + "step_preserves_invariant", + Verdict.CONFIRMED, + ), + Claim( + "Decrementing at zero keeps the counter at zero", + "dec_at_zero_keeps_zero", + Verdict.CONFIRMED, + ), + Claim( + "The counter value is always non-negative", + "counter_non_neg_alt", + Verdict.DISPUTED, + "tautology: `m = m` proves nothing about non-negativity", + ), + Claim( + "The counter value is always non-negative", + "counter_non_neg_large", + Verdict.DISPUTED, + "vacuous precondition: `100 < m` narrows scope, doesn't cover m = 0", + ), + Claim( + "The counter value is always non-negative", + "counter_lower_bound", + Verdict.DISPUTED, + "weakened postcondition: concludes `-1 <= m` instead of `0 <= m`", + ), +) + +CANON_CLAIMS: tuple[Claim, ...] = ( + # Upstream's `canon`. Two conclusions are opaque predicates + # (`AllConstraintsValid`, `AllEdgesValid`); `remove_node_cleans_up` is + # the three-conjunct conclusion whose first conjunct alone is + # `remove_node_drops_id`. + Claim( + "All constraint targets reference existing nodes", + "constraint_targets_exist", + Verdict.CONFIRMED, + ), + Claim( + "All edge endpoints reference existing nodes", + "edge_endpoints_exist", + Verdict.CONFIRMED, + ), + Claim( + "Adding a node with an existing ID is a no-op", + "add_existing_node_is_noop", + Verdict.CONFIRMED, + ), + Claim( + "Removing a node cleans up related constraints and edges", + "remove_node_cleans_up", + Verdict.CONFIRMED, + ), + Claim( + "Removing a node cleans up related constraints and edges", + "remove_node_drops_id", + Verdict.DISPUTED, + "missing conjunct: only checks the node is removed, doesn't verify " + "constraint/edge cleanup", + ), + Claim( + "All constraint targets reference existing nodes", + "constraint_targets_exist_empty", + Verdict.DISPUTED, + "vacuous precondition: requiring no constraints makes the conclusion trivially " + "true", + ), +) + +COLORWHEEL_CLAIMS: tuple[Claim, ...] = ( + # Upstream's `colorwheel`. The domain where the weakest arm falls apart + # upstream, and the split is legible: it confirms the two theorems whose + # conclusions are visible arithmetic (`always_five_colors`, + # `contrast_pair_indices_valid`) and disputes the four whose conclusions + # are named predicates it has never seen unfolded. + Claim( + "The base hue is always in valid range", + "base_hue_in_range", + Verdict.CONFIRMED, + ), + Claim( + "There are always exactly 5 colors in the palette", + "always_five_colors", + Verdict.CONFIRMED, + ), + Claim( + "Every color has valid saturation and lightness values", + "all_colors_valid", + Verdict.CONFIRMED, + ), + Claim( + "Contrast pair indices are valid (between 0 and 4)", + "contrast_pair_indices_valid", + Verdict.CONFIRMED, + ), + Claim( + "When a mood is set (not Custom), all colors satisfy the mood constraints", + "mood_constraints_satisfied", + Verdict.CONFIRMED, + ), + Claim( + "Hues follow the selected harmony pattern", + "hues_follow_harmony", + Verdict.CONFIRMED, + ), + Claim( + "There are always exactly 5 colors in the palette", + "palette_non_empty", + Verdict.DISPUTED, + "weakened postcondition: concludes the palette is non-empty instead of exactly " + "5", + ), +) + +DELEGATION_CLAIMS: tuple[Claim, ...] = ( + # Upstream's `delegation-auth`. `delegate_non_existent_is_noop` is one of + # only two items every single-call arm gets wrong on every model upstream + # tried: the requirement reads as 'both subjects missing' and the + # hypothesis says 'at least one missing', so the theorem is *stronger* + # than what was asked -- which still counts as expressing it. + Claim( + "All granted capabilities reference existing subjects", + "grant_subjects_exist", + Verdict.CONFIRMED, + ), + Claim( + "Delegation endpoints (from, to) must be existing subjects", + "delegation_endpoints_exist", + Verdict.CONFIRMED, + ), + Claim( + "Edge IDs are always less than the next allocator (freshness)", + "edge_ids_fresh", + Verdict.CONFIRMED, + ), + Claim( + "Granting a capability to a non-existent subject is a no-op", + "grant_non_existent_is_noop", + Verdict.CONFIRMED, + ), + Claim( + "Delegating between non-existent subjects is a no-op", + "delegate_non_existent_is_noop", + Verdict.CONFIRMED, + ), + Claim( + "Revoking a non-existent delegation is a no-op", + "revoke_non_existent_is_noop", + Verdict.CONFIRMED, + ), + Claim( + "Granting a capability to a non-existent subject is a no-op", + "grant_non_existent_is_noop_init", + Verdict.DISPUTED, + "vacuous precondition: `m = Init` restricts the claim to the empty policy only", + ), +) + +KANBAN_CLAIMS: tuple[Claim, ...] = ( + # Upstream's `kanban`, and the sharpest pair in the benchmark: + # `no_card_duplicates` and `card_partition_no_dups` have identical + # statements. Which one is faithful depends entirely on the requirement + # it is set against, so no amount of reading the Lean decides it. + Claim( + "Column names are unique (no duplicate columns)", + "columns_are_unique", + Verdict.CONFIRMED, + ), + Claim( + "Every card appears in exactly one column (exact partition)", + "card_in_exactly_one_column", + Verdict.CONFIRMED, + ), + Claim( + "No card ID appears twice across all lanes (no duplicates)", + "no_card_duplicates", + Verdict.CONFIRMED, + ), + Claim( + "Each column respects its WIP limit (number of cards does not exceed the " + "limit)", + "wip_limits_respected", + Verdict.CONFIRMED, + ), + Claim( + "Adding a card to a full column is a no-op", + "add_card_to_full_column_is_noop", + Verdict.CONFIRMED, + ), + Claim( + "The card allocator is always fresh (no allocated ID reused)", + "allocator_always_fresh", + Verdict.CONFIRMED, + ), + Claim( + "Lanes and WIP maps are defined exactly for existing columns", + "lanes_and_wip_match_columns", + Verdict.CONFIRMED, + ), + Claim( + "Moving a card preserves the total number of cards", + "move_card_preserves_total", + Verdict.DISPUTED, + "tautology: compares one expression to itself", + ), + Claim( + "Every card appears in exactly one column (exact partition)", + "card_partition_no_dups", + Verdict.DISPUTED, + "missing conjunct: proves only that IDs are distinct, not the bidirectional " + "membership that makes it a partition", + ), +) + + +@dataclasses.dataclass(frozen=True) +class Domain: + """One body of Lean and the claims made about it.""" + + name: str + corpus: str + claims: tuple[Claim, ...] + + +DOMAINS: dict[str, Domain] = { + "counter": Domain("counter", COUNTER_CORPUS, COUNTER_CLAIMS), + "canon": Domain("canon", CANON_CORPUS, CANON_CLAIMS), + "colorwheel": Domain("colorwheel", COLORWHEEL_CORPUS, COLORWHEEL_CLAIMS), + "delegation": Domain("delegation", DELEGATION_CORPUS, DELEGATION_CLAIMS), + "kanban": Domain("kanban", KANBAN_CORPUS, KANBAN_CLAIMS), +} + + +# --------------------------------------------------------------------------- +# Driving one claim through a mode. Each claim gets its own agent instances, so +# nothing an audit learns can leak into the next one through a shared history. +# --------------------------------------------------------------------------- + + +class Mode(enum.StrEnum): + TWO_PASS = "two-pass" + NAIVE = "naive" + + +@dataclasses.dataclass(frozen=True) +class Audit: + """One claim's result: what the pipeline decided, and what it read on the way. + + ``comparison`` is None when the pipeline never produced a well-formed verdict + -- the model kept emitting an incoherent one and `RetryLLMHandler` ran out of + attempts. That is upstream's third status, ``error``: not a confirmation and + not a dispute, and it counts against the run rather than being dropped. + """ + + domain: str + claim: Claim + statement: str + verdict: Verdict | None + explanation: str + comparison: Comparison | None # None in naive mode, which has no taxonomy + back_translation: Informalization | None # None outside two-pass mode + error: str | None = None + + @property + def correct(self) -> bool: + return self.verdict is self.claim.expected + + @property + def label(self) -> str: + return f"{self.domain}/{self.claim.theorem}" + + +def audit_claim(domain: Domain, claim: Claim, mode: Mode) -> Audit: + """Audit one requirement/theorem pair under the given mode. + + A claim the model cannot produce a decodable verdict for becomes an ``error`` + result rather than an exception: one intractable item should cost one item, + not the other thirty-nine. + """ + statement = statement_of(domain.corpus, claim.theorem) + back: Informalization | None = None + comparison: Comparison | None = None + try: + if mode is Mode.NAIVE: + # The floor: a yes/no and a sentence, no taxonomy. + judgement = NaiveAuditor().audit(claim.requirement, statement) + match, explanation = judgement.match, judgement.explanation + else: + # Pass 1 receives `statement`. There is nowhere in this call for + # `claim.requirement` to go. + back = Informalizer().informalize(statement) + comparison = Comparator().compare(claim.requirement, statement, back) + match, explanation = comparison.match, comparison.explanation + except Exception as exc: + return Audit( + domain.name, + claim, + statement, + None, + "", + comparison, + back, + f"{type(exc).__name__}: {exc}", + ) + verdict = Verdict.CONFIRMED if match else Verdict.DISPUTED + return Audit(domain.name, claim, statement, verdict, explanation, comparison, back) + + +async def audit_all(domains: typing.Sequence[Domain], mode: Mode) -> list[Audit]: + """Audit every claim in every domain concurrently -- independent by + construction, since each gets its own agent instances.""" + return list( + await asyncio.gather( + *( + asyncio.to_thread(audit_claim, domain, claim, mode) + for domain in domains + for claim in domain.claims + ) + ) + ) + + +# --------------------------------------------------------------------------- +# The deterministic half. The blog's between-pass checks are diagnostics over +# typed values, so they are a loop in Python rather than another model call: +# code does what code can decide, and the model is asked only what needs +# judgment. +# --------------------------------------------------------------------------- + + +def pre_checks(audits: typing.Sequence[Audit]) -> list[str]: + """Flag back-translations rated trivial, and distinct requirements whose + theorems were read as guaranteeing the same thing.""" + notes: list[str] = [] + seen: dict[str, Claim] = {} + for audit in audits: + if (back := audit.back_translation) is None: + continue + if back.strength is Strength.TRIVIAL: + notes.append( + f"{audit.label} was read as a trivial claim ({back.conclusion})" + ) + key = f"{audit.domain}: {' '.join(back.conclusion.lower().split())}" + if (earlier := seen.get(key)) is not None: + if earlier.requirement != audit.claim.requirement: + notes.append( + f"{audit.label} and {earlier.theorem} were read as " + "guaranteeing the same thing, but formalize different " + "requirements" + ) + else: + seen[key] = audit.claim + return notes + + +# --------------------------------------------------------------------------- +# Reporting. The two error directions are reported apart: a missed dispute is a +# weak theorem waved through, which is the failure ClaimCheck exists to prevent, +# while a false dispute costs a developer an argument with the tool. +# --------------------------------------------------------------------------- + + +def report(audits: typing.Sequence[Audit], mode: Mode) -> None: + print(f"\n{'=' * 78}\nClaimCheck audit -- strategy: {mode.value}\n{'=' * 78}\n") + + for audit in audits: + mark = "ok " if audit.correct else "MISS" + category = ( + f" [{audit.comparison.weakening.value}]" + if audit.comparison is not None + and audit.comparison.weakening is not Weakening.NONE + else "" + ) + decided = "error" if audit.verdict is None else audit.verdict.value + print(f"[{mark}] {audit.label}: {decided}{category}") + print(f" requirement: {audit.claim.requirement}") + print(f" statement: {' '.join(audit.statement.split())}") + if audit.back_translation is not None: + back = audit.back_translation + print( + f" read as: {back.natural_language} " + f"(strength: {back.strength.value})" + ) + if audit.comparison is not None and audit.comparison.discrepancy: + print(f" discrepancy: {audit.comparison.discrepancy}") + elif audit.explanation: + print(f" reasoning: {audit.explanation}") + if audit.error: + print(f" no verdict: {audit.error}") + if not audit.correct: + print(f" EXPECTED {audit.claim.expected.value}: {audit.claim.why}") + print() + + if notes := pre_checks(audits): + print("Pre-check diagnostics (deterministic, no model involved):") + for note in notes: + print(f" - {note}") + print() + + # A "missed dispute" is an unfaithful theorem the audit confirmed. Claims the + # pipeline never returned a verdict for are counted apart from both error + # directions -- they are a failure of the harness, not of judgment. + errored = [a for a in audits if a.verdict is None] + missed = [ + a + for a in audits + if a.verdict is not None + and a.claim.expected is Verdict.DISPUTED + and not a.correct + ] + false_alarms = [ + a + for a in audits + if a.verdict is not None + and a.claim.expected is Verdict.CONFIRMED + and not a.correct + ] + correct = sum(a.correct for a in audits) + + # Per-domain accuracy as well as overall: the domains differ in difficulty, + # and an aggregate hides which one a strategy actually struggles with. + by_domain: dict[str, list[Audit]] = {} + for audit in audits: + by_domain.setdefault(audit.domain, []).append(audit) + if len(by_domain) > 1: + for name, group in by_domain.items(): + hits = sum(a.correct for a in group) + print(f" {name:12} {hits}/{len(group)} ({hits / len(group):.1%})") + + print( + f"Accuracy: {correct}/{len(audits)} " + f"({correct / len(audits):.1%})\n" + f" unfaithful theorems waved through: {len(missed)}" + + (f" ({', '.join(a.label for a in missed)})" if missed else "") + + f"\n faithful theorems disputed: {len(false_alarms)}" + + (f" ({', '.join(a.label for a in false_alarms)})" if false_alarms else "") + + f"\n no verdict (retries exhausted): {len(errored)}" + + (f" ({', '.join(a.label for a in errored)})" if errored else "") + ) + + +# --------------------------------------------------------------------------- +# The premise, checked. ClaimCheck is only interesting if the formal artifacts +# really are proved -- otherwise a disputed theorem might just be a broken one. +# `formalization.py` (LEAP) already drives a real Lean 4 + Mathlib toolchain, so +# the check reuses its kernel rather than restating it. Imported inside the +# function, as `world_model_agent.py` imports `gridworlds`, so the example has no +# Lean dependency unless the check is asked for. +# --------------------------------------------------------------------------- + + +def verify_corpus(domains: typing.Sequence[Domain]) -> bool: + """Compile each domain's corpus with Lean, and report what it proves.""" + # The examples are importable as ``docs.source.llm_examples...`` from the + # repository root, which is on ``sys.path`` under the harness but not when + # this file is run directly; add it so both invocations work. + sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[4])) + from docs.source.llm_examples.autoformalization.formalization import ( + _SORRY, + LeanKernel, + ) + + kernel = LeanKernel() + if not kernel.available(): + print( + f"Lean project not built at {kernel.project!r}; skipping verification.\n" + "Build it once (see formalization.py --check-toolchain):\n" + " elan default stable\n" + f" cd {kernel.project} && lake exe cache get && lake build" + ) + return False + + total = 0 + for domain in domains: + theorems = re.findall(r"^theorem (\w+)", domain.corpus, re.MULTILINE) + print( + f"Compiling {domain.name} ({len(theorems)} theorems) with " + "Lean 4 + Mathlib ..." + ) + result = kernel.compile(domain.corpus) + if not result.ok: + raise SystemExit( + f"The {domain.name} corpus does not compile:\n{result.messages}" + ) + if _SORRY.search(domain.corpus): + raise SystemExit( + f"The {domain.name} corpus contains `sorry`; it is not proved." + ) + total += len(theorems) + print( + f"VERIFIED: {total} theorems, 0 errors, no `sorry`. Every claim below is " + "proved.\nThe audit that follows is not about whether they are true.\n" + ) + return True + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + # Named `--strategy`, not `--mode`: the harness parses its own flags with + # argparse's prefix matching on, so a script flag named `--mode` is swallowed + # as an abbreviation of the harness's `--model`. + parser.add_argument( + "--strategy", + type=Mode, + choices=list(Mode), + default=Mode.TWO_PASS, + help="Audit strategy: the two-pass split, in which the informalizer " + "never sees the requirement, or the naive floor -- one call, " + "'does this match?'", + ) + parser.add_argument( + "--verify", + action="store_true", + help="Compile the corpus with a real Lean toolchain first, establishing " + "that every theorem audited below is actually proved", + ) + parser.add_argument( + "--verify-only", + action="store_true", + help="Compile the corpus and exit, without calling any model", + ) + parser.add_argument( + "--domain", + choices=[*DOMAINS, "all"], + default="all", + help="Which of upstream's five benchmark domains to audit, or all of them", + ) + parser.add_argument( + "--limit", + type=int, + default=None, + help="Audit only the first N claims of each domain (a cheap smoke test)", + ) + args = parser.parse_args() + + domains = list(DOMAINS.values()) if args.domain == "all" else [DOMAINS[args.domain]] + if args.limit: + domains = [ + dataclasses.replace(d, claims=d.claims[: args.limit]) for d in domains + ] + + if args.verify_only: + verify_corpus(domains) + return + if args.verify: + verify_corpus(domains) + + report(asyncio.run(audit_all(domains, args.strategy)), args.strategy) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoformalization/auditing_agents.py b/docs/source/llm_examples/autoformalization/auditing_agents.py new file mode 100644 index 000000000..52b8399bb --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing_agents.py @@ -0,0 +1,247 @@ +"""Agents and model-boundary types for the ClaimCheck audit in `auditing.py`. + +Two agents read Lean theorem statements and judge whether they express a +natural-language requirement, plus the typed values that cross the model +boundary. + +**This module is deliberately small, and that is the point.** The harness builds +a template's system prompt partly from the source of the module the template is +defined in (see the prompt-assembly table in +`effectful.handlers.llm.types.Template`), so everything sharing a file with an +`Agent` is shown to it verbatim. Keeping these two agents in a module of their +own is what makes it true that the informalizer sees a Lean statement and nothing +else. Nothing here knows anything about the material being audited; `auditing.py` +imports this module and is never imported by it. The two single-call ablations +live in `auditing_single` and `auditing_naive` for the same reason: their prompts +differ from these, and a shared module would show each arm the others'. +""" + +import dataclasses +import enum + +import pydantic.dataclasses + +from effectful.handlers.llm import Agent, Template + + +class Verdict(enum.StrEnum): + CONFIRMED = "confirmed" + DISPUTED = "disputed" + + +class Weakening(enum.StrEnum): + """How a theorem can fail to mean its requirement -- the blog's taxonomy.""" + + NONE = "none" + TAUTOLOGY = "tautology" + WEAKENED_CONCLUSION = "weakened-conclusion" + NARROWED_SCOPE = "narrowed-scope" + MISSING_CASE = "missing-case" + WRONG_PROPERTY = "wrong-property" + + +# --------------------------------------------------------------------------- +# Types crossing the model boundary. A field's ``metadata={"description": ...}`` +# is inlined by pydantic into that field's JSON schema and rendered into the +# system prompt, so per-field guidance reaches the model through the type and no +# prompt has to restate it. +# --------------------------------------------------------------------------- + + +class Strength(enum.StrEnum): + TRIVIAL = "trivial" + WEAK = "weak" + MODERATE = "moderate" + STRONG = "strong" + + +@pydantic.dataclasses.dataclass(frozen=True) +class Informalization: + """Pass 1's output: what a Lean statement says, read on its own terms. + + Produced without sight of the requirement the theorem was written for, which + is the whole mechanism -- a back-translation that agreed with the requirement + because it had been shown the requirement would be worth nothing. + """ + + natural_language: str = dataclasses.field( + metadata={ + "description": "One sentence of plain English for what this theorem " + "guarantees. Be literal: describe what the statement says, not what " + "you suppose its author was aiming at." + } + ) + hypotheses: str = dataclasses.field( + metadata={ + "description": "What must hold for the guarantee to apply, in English; " + "'none' if the statement holds unconditionally." + } + ) + conclusion: str = dataclasses.field( + metadata={"description": "What is guaranteed, in English."} + ) + scope: str = dataclasses.field( + metadata={ + "description": "What the guarantee ranges over: every state of the " + "system, one particular state, states satisfying some restriction, " + "etc." + } + ) + strength: Strength = dataclasses.field( + metadata={ + "description": "'trivial' if the conclusion restates a hypothesis or " + "holds for every value of the types involved regardless (e.g. a " + "natural number being non-negative, or both sides of an equation " + "being the same term); 'weak' if it says very little; 'moderate' if " + "it is a substantive claim; 'strong' if it constrains behaviour " + "sharply." + } + ) + confidence: float = dataclasses.field( + metadata={"description": "0-1, how sure you are this reading is faithful."} + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Comparison: + """Pass 2's verdict on one requirement/theorem pair. + + ``__post_init__`` certifies the verdict is internally coherent before it is + ever returned: a match is exactly a `Weakening.NONE`, and a mismatch has to + say what is wrong. An incoherent answer raises, and `RetryLLMHandler` hands + the message back to the model as the next turn -- so "matches, but it's a + tautology" is not a verdict this pipeline can emit. + """ + + match: bool = dataclasses.field( + metadata={ + "description": "True only if the theorem expresses the whole of the " + "requirement. A theorem that is stronger than the requirement still " + "matches; one that is weaker, narrower, or about something else " + "does not." + } + ) + weakening: Weakening = dataclasses.field( + metadata={ + "description": "The category of divergence: 'none' when and only when " + "match is true." + } + ) + discrepancy: str = dataclasses.field( + metadata={ + "description": "What the requirement asks for that the theorem does " + "not deliver. Empty when match is true." + } + ) + explanation: str = dataclasses.field( + metadata={"description": "Brief reasoning for the verdict."} + ) + + def __post_init__(self) -> None: + if self.match and self.weakening is not Weakening.NONE: + raise ValueError( + f"incoherent verdict: match is true but weakening is " + f"{self.weakening.value!r}. If the theorem really expresses the " + "requirement the weakening is 'none'; otherwise match is false." + ) + if not self.match and self.weakening is Weakening.NONE: + raise ValueError( + "incoherent verdict: match is false but weakening is 'none'. " + "Name the category of the divergence." + ) + if not self.match and not self.discrepancy.strip(): + raise ValueError( + "match is false but no discrepancy is given; say what the " + "requirement asks for that the theorem does not deliver." + ) + + @property + def verdict(self) -> Verdict: + return Verdict.CONFIRMED if self.match else Verdict.DISPUTED + + +# --------------------------------------------------------------------------- +# Pass 1. `informalize` takes a Lean statement and nothing else: there is no +# parameter for a requirement, this Agent's history contains no turn in which +# one appeared, and no Tool in scope can go and find one. That is the entire +# separation -- not an instruction the model is trusted to obey. +# --------------------------------------------------------------------------- + + +class Informalizer(Agent): + """You read Lean 4 theorem statements and say, in plain English, exactly what + they guarantee. You are a translator, not a sympathetic reader: you report + what the statement says, never what you imagine it was for. You are not shown + why any theorem was written, and you should not speculate about it.""" + + @Template.define + def informalize(self, statement: str) -> Informalization: + """Translate this Lean 4 theorem statement into English, as literally as + you can. + + ```lean + {statement} + ``` + + Separate what is assumed (the hypotheses) from what is guaranteed (the + conclusion), and say what the guarantee ranges over. Then rate how much + the statement actually claims -- be blunt about this. A conclusion that + holds for every value of the types involved, or that merely repeats a + hypothesis, is trivial no matter how substantial the theorem's name + makes it sound. + + Read only the statement in front of you. Do not guess at intent. + """ + + +# --------------------------------------------------------------------------- +# Pass 2. This agent does see the requirement -- comparing is its job. What it +# gets from pass 1 is a reading of the formal statement produced in ignorance of +# that requirement, so agreement between them is evidence. +# --------------------------------------------------------------------------- + + +class Comparator(Agent): + """You check whether a formal theorem carries the weight a natural-language + requirement puts on it. You assume the proof is correct: you are not auditing + the proof, you are auditing the claim. You are strict -- a theorem that is + true, proved, and beside the point is a finding -- but not pedantic about + wording, since only the meaning has to survive.""" + + @Template.define + def compare( + self, requirement: str, statement: str, back_translation: Informalization + ) -> Comparison: + """Decide whether this theorem expresses this requirement. + + **Requirement, as written by the person who asked for it:** + {requirement} + + **The theorem said to formalize it:** + ```lean + {statement} + ``` + + **Back-translation** -- what the statement says, according to a reader + who was shown the statement alone and never saw the requirement above: + {back_translation} + + Watch for the ways a proved theorem can still miss: + + 1. **tautology** -- the conclusion restates a hypothesis, or holds for + every value of the types involved, so nothing is established. + 2. **weakened-conclusion** -- the theorem guarantees less than was asked + (a looser bound, a weaker relation). + 3. **narrowed-scope** -- the theorem only covers a subset of the + cases the requirement describes. + 4. **missing-case** -- the requirement asks for several things and the + theorem delivers some of them. + 5. **wrong-property** -- the theorem is about something else, however + adjacent. + + A theorem *stronger* than the requirement still matches; do not flag + rephrasing. But if the back-translation rates the statement trivial, the + requirement had better be trivial too. Judge the statement, not its name: + a theorem called after the property it was meant to prove is no evidence + that it proves it. + """ diff --git a/docs/source/llm_examples/autoformalization/auditing_naive.py b/docs/source/llm_examples/autoformalization/auditing_naive.py new file mode 100644 index 000000000..890a828bf --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing_naive.py @@ -0,0 +1,68 @@ +"""The naive audit arm: one call, a verdict, and a sentence. + +Alone in a module because a template's system prompt includes the source of its +defining module, so agents sharing a file are shown each other's prompts and +types. The arms differ, so they do not share a file. +""" + +import dataclasses +import enum + +import pydantic.dataclasses + +from effectful.handlers.llm import Agent, Template + + +class NaiveVerdict(enum.StrEnum): + JUSTIFIED = "JUSTIFIED" + NOT_JUSTIFIED = "NOT_JUSTIFIED" + + +@pydantic.dataclasses.dataclass(frozen=True) +class NaiveJudgement: + """A verdict and a sentence of justification.""" + + verdict: NaiveVerdict = dataclasses.field( + metadata={ + "description": "JUSTIFIED if the theorem captures the requirement, " + "NOT_JUSTIFIED if there is a meaningful discrepancy." + } + ) + explanation: str = dataclasses.field( + metadata={"description": "Brief explanation of your verdict."} + ) + + @property + def match(self) -> bool: + return self.verdict is NaiveVerdict.JUSTIFIED + + +class NaiveAuditor(Agent): + """You check whether verified Lean theorems correctly formalize the natural + language requirements they are said to capture.""" + + @Template.define + def audit(self, requirement: str, statement: str) -> NaiveJudgement: + """Does this Lean theorem faithfully capture the requirement below? + + ## Natural Language Requirement + + > {requirement} + + ## Lean Theorem + + ```lean + {statement} + ``` + + ## Instructions + + - **JUSTIFIED** if the theorem's statement expresses the requirement (it + may be stronger, that's fine). + - **NOT_JUSTIFIED** if there is a meaningful discrepancy: the theorem is + weaker, proves something different, is vacuous, or misses key aspects. + + Invariant hypotheses (e.g. ``Inv m``) are expected and normal -- don't + count them as discrepancies. A theorem that extracts a concrete + consequence from an invariant is useful, not vacuous. + """ diff --git a/docs/source/llm_examples/autoformalization/formalization.py b/docs/source/llm_examples/autoformalization/formalization.py new file mode 100644 index 000000000..6f667cc16 --- /dev/null +++ b/docs/source/llm_examples/autoformalization/formalization.py @@ -0,0 +1,961 @@ +"""LEAP: blueprint-driven formal theorem proving over a *real* Lean compiler. + +Implements the core of "LEAP: Supercharging LLMs for Formal Mathematics with +Agentic Frameworks" (arXiv:2606.03303). The paper's diagnosis is that general +LLMs reason well informally but "struggle to generate mechanically verifiable +proofs in formal languages like Lean" -- one-shot formalization of a hard theorem +essentially never compiles. Its fix is to treat proving as an *orchestration* +problem: register the theorem as the root of an AND-OR DAG, attempt a *direct* +proof with compiler-feedback revision, and on failure *decompose* it -- draft an +informal blueprint proposing intermediate lemmas, translate that into a Lean +*sketch* that proves the goal assuming the lemmas (``sorry`` placeholders), have an +LLM reviewer judge the decomposition, and recurse on the subgoals -- reusing proved +lemmas across branches via hierarchical memoization. + +Unlike the sibling examples, which fake their environment (a static in-memory +index instead of live web search), LEAP's environment *is* the load-bearing part, +so we do not fake it: the ``VERIFIER (LEAN)`` of the paper's Figure 1 is a real +Lean 4 + Mathlib toolchain, invoked as a subprocess. A proof that does not compile +raises with the actual Lean error, and the harness's ``RetryLLMHandler`` feeds that +error back -- the paper's "continuous interaction with the Lean compiler" is a real +compile loop, not a simulation. Each of the paper's named components falls out of +an ordinary effectful idiom: + + * Grounded proofs by construction. A ``LeanProof`` certifies *at decode time* + that `` := by `` compiles under Lean with no errors and no + ``sorry`` -- the same decode-time certification ``scientist_one.py`` uses for + citations, except the ground truth is a theorem prover rather than an index. + An uncompilable proof is not a well-typed ``LeanProof``; it raises, and + ``RetryLLMHandler`` feeds the compiler diagnostic back (the paper's ``REVISER``). + + * The sketch is the same certification with a richer preamble. A decomposition's + sketch proves the goal *assuming* its proposed lemmas: the search installs the + lemmas as ``sorry`` stubs in the compile preamble, so the sketch's own tactics + must be ``sorry``-free (checked) while depending on the sorried lemmas -- exactly + the paper's "main theorem body is ``sorry``-free, ``sorry`` permitted in the + proposed lemma statements". + + * Interleaved informal->formal planning. Both paths pass through an informal + step before Lean: the ``NLProver`` writes an informal argument the + ``FormalProver`` formalizes, and the ``BlueprintAgent`` drafts an informal + decomposition the ``SketchAgent`` turns into a Lean sketch (the two-stream + shape of ``scholar_peer.py``). + + * Tools scoped by class: only the formalizing agents subclass ``LeanAgent`` and + hold the ``check`` tool that compiles a candidate against the live goal state + and returns Lean's messages -- the compiler-in-the-loop. The planning and + reviewing agents are closed-book by construction, no "do not compile" + instruction needed (the encapsulation idiom of ``scholar_peer.py``). + + * Verification-guided proof search. Compiler verification is necessary but not + sufficient: a sketch can compile while introducing a subgoal no simpler than + its parent (paper Figure 3). The ``Reviewer`` LLM acts as a search filter that + rejects such decompositions, and the ``state_writer`` refuses any subgoal that + would reintroduce an ancestor -- preserving the DAG's acyclicity. Search is a + DFS with backtracking over blueprints. + + * Hierarchical memoization via the AND-OR DAG. Goals are OR nodes keyed by their + (normalized) statement; a decomposition is an AND node whose parent is proved + once all its child subgoals are. A lemma proved in one branch is stored as a + real Lean declaration and (a) reused verbatim if the same statement resurfaces + in another branch -- turning a would-be decomposition into a direct proof -- + and (b) carried in every downstream compile preamble, so the final assembled + proof of the root is one real Lean file that compiles end-to-end with no + ``sorry``. + +Demonstrates: +- Decode-time certification against a *real external tool* (the Lean compiler), + so ``RetryLLMHandler`` turns an uncompilable proof into a compiler-feedback + revision -- the certification idiom of ``scientist_one.py`` with a prover as + ground truth +- A ContextVar carrying per-goal compile state (preamble + goal), read ambiently + by ``LeanProof.__post_init__`` and the ``check`` tool, scoped to the pipeline + (the ``WORKSPACE``/``CUTOFF`` idiom of ``scientist_one``/``paper_orchestra``) +- A class-scoped compiler tool offered to the formalizing agents via the Agent MRO + and invisible to the closed-book planning/review agents (``scholar_peer.py``) +- An AND-OR DAG with hierarchical memoization, DFS backtracking, an LLM reviewer + as a search filter, and a state-writer acyclicity guard -- the paper's Figure 1 +- End-to-end verification: the assembled proof tree is emitted as one Lean file and + compiled with no ``sorry``, the way ``scientist_one``'s audit re-derives its + evidence +""" + +# Simplifications vs. the source: +# - No Lean-IMO-Bench / Putnam. The paper proves olympiad-level theorems; this +# composes a proof of a small, self-contained target so the example runs in +# minutes, not a leaderboard. The architecture -- direct-then-decompose over an +# AND-OR DAG with memoization -- is the same. +# - LeanSearch is a compile loop, not premise retrieval. The paper retrieves premises +# with LeanSearch; here the ``check`` tool compiles a candidate against the live +# goal and returns Lean's messages (errors / remaining goals), which is the +# compiler-interaction half of that loop. Mathlib's own ``exact?``/``apply?`` remain +# available to the model *inside* a proof, so premise search still happens -- in Lean. +# - One reviewer pass, single-vote. The decomposition reviewer judges once rather +# than by majority vote (contrast ``scientist_one``'s majority-vote audit); the +# acyclicity guard is deterministic Python. +# - Memoization is textual. Two lemmas are "the same" node when their normalized +# statements match textually (whitespace-collapsed), not up to Lean-level +# defeq/alpha -- enough to share the reusable-lemma story without an elaboration +# check on every pair. + +import argparse +import collections.abc +import contextvars +import dataclasses +import hashlib +import os +import re +import shutil +import subprocess +import textwrap + +import pydantic.dataclasses + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# The Lean compiler -- the ground truth every proof is certified against. This is +# the paper's ``VERIFIER (LEAN)``: a real Lean 4 + Mathlib toolchain shelled out to, +# not a stand-in. A proof is valid iff Lean accepts the file with no error message. +# --------------------------------------------------------------------------- + +# Where the Mathlib lake project lives. Built once (elan + `lake exe cache get`); +# see this module's header. Override with LEAP_LEAN_PROJECT. +LEAN_PROJECT = os.environ.get( + "LEAP_LEAN_PROJECT", os.path.expanduser("~/.cache/leap-lean/leapproj") +) +# Every compiled fragment opens with this; `import Mathlib` pulls the whole library +# so the model may use any tactic/lemma it knows (`ring`, `omega`, `simp`, `exact?`). +PRELUDE = "import Mathlib\n" + + +@dataclasses.dataclass(frozen=True) +class LeanResult: + """The outcome of compiling a Lean fragment: ``ok`` is true iff Lean reported no + error (``sorry`` warnings are not errors). ``messages`` is Lean's stdout+stderr, + fed back to the model verbatim on failure -- the raw compiler diagnostic.""" + + ok: bool + messages: str + + +def _lake_bin() -> str: + """Locate the ``lake`` executable, tolerating a not-yet-on-PATH elan install.""" + for cand in ( + os.environ.get("LAKE"), + shutil.which("lake"), + os.path.expanduser("~/.elan/bin/lake"), + ): + if cand and os.path.exists(cand): + return cand + return "lake" + + +class LeanKernel: + """Compiles Lean source via ``lake env lean`` in the Mathlib project, with an + in-memory cache keyed by source text so identical fragments (retries, repeated + tool calls, the same proved lemma seen twice) compile at most once. Importing + all of Mathlib per check is slow; the cache is what keeps the search tractable.""" + + def __init__(self, project: str = LEAN_PROJECT, timeout: float = 120.0) -> None: + self.project = project + self.timeout = timeout + self._cache: dict[str, LeanResult] = {} + + def available(self) -> bool: + return os.path.isdir(os.path.join(self.project, ".lake")) + + def compile(self, source: str) -> LeanResult: + """Compile a full Lean source string and return the result (cached).""" + key = hashlib.sha256(source.encode()).hexdigest() + if key in self._cache: + return self._cache[key] + env = dict(os.environ) + env["PATH"] = ( + os.path.expanduser("~/.elan/bin") + os.pathsep + env.get("PATH", "") + ) + # A scratch file inside the project's build dir so `lake env` resolves imports. + scratch = os.path.join(self.project, f".leap_scratch_{key[:12]}.lean") + try: + with open(scratch, "w") as fh: + fh.write(source) + proc = subprocess.run( + [_lake_bin(), "env", "lean", scratch], + cwd=self.project, + capture_output=True, + text=True, + timeout=self.timeout, + env=env, + ) + out = (proc.stdout + proc.stderr).strip() + # `lean` exits non-zero on error; `sorry` and linter notes are warnings. + ok = proc.returncode == 0 and "error:" not in out + except subprocess.TimeoutExpired: + ok, out = False, f"Lean timed out after {self.timeout}s (proof too slow)." + finally: + if os.path.exists(scratch): + os.remove(scratch) + result = LeanResult(ok, out or ("no output" if ok else "unknown error")) + self._cache[key] = result + return result + + +# The compile context for the goal currently being worked. ``LeanProof`` and the +# ``check`` tool read it ambiently -- through a ContextVar rather than a bare global, +# so it is scoped to the pipeline and safe if goals are ever worked concurrently. +# Exactly ``scientist_one``'s WORKSPACE / ``paper_orchestra``'s CUTOFF pattern. +@dataclasses.dataclass(frozen=True) +class LeanContext: + kernel: LeanKernel + preamble: str # PRELUDE + proved lemmas + (for a sketch) the sorry-stub lemmas + decl: ( + str # the goal declaration header, e.g. "theorem leap_goal (n : ℕ) : n + 0 = n" + ) + + +LEAN_CTX: contextvars.ContextVar[LeanContext] = contextvars.ContextVar("LEAN_CTX") + +# A proof body may not smuggle in `sorry` (or its cousins): the main goal must be +# genuinely closed. `sorry` is legitimate only in the search-generated lemma stubs, +# which live in the preamble, never in model-authored tactics. +_SORRY = re.compile(r"\b(sorry|admit|sorryAx)\b") + + +def assemble(decl: str, tactics: str, preamble: str) -> str: + """Build the full Lean source for ``decl := by `` under ``preamble``.""" + body = textwrap.indent(tactics.strip(), " ") + return f"{preamble}\n\n{decl} := by\n{body}\n" + + +def with_ctx[T](ctx: "LeanContext", fn: collections.abc.Callable[[], T]) -> T: + """Run ``fn`` with ``LEAN_CTX`` bound to ``ctx`` for exactly that call, so the + ``check`` tool and the decode-time certifications read the right goal/preamble. + One balanced set/reset per call -- no fragile nesting across a whole loop body.""" + token = LEAN_CTX.set(ctx) + try: + return fn() + finally: + LEAN_CTX.reset(token) + + +# --------------------------------------------------------------------------- +# Types crossing the model boundary +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class LeanProof: + """A tactic-block proof of the goal currently in scope, certified at decode time. + + ``__post_init__`` assembles `` := by `` under the in-scope + preamble and compiles it with the real Lean kernel; an uncompilable proof (or one + that tries to use ``sorry``) raises, and ``RetryLLMHandler`` feeds Lean's own + error message back so the model revises against the compiler. Used for both the + direct proof and the decomposition sketch -- they differ only in the preamble the + search installs (a sketch's preamble carries the proposed lemmas as ``sorry`` + stubs, so the goal may lean on them while its own tactics stay ``sorry``-free).""" + + tactics: str = dataclasses.field( + metadata={ + "description": "The tactic block that proves the goal, i.e. what follows " + "`:= by`. Do not include the theorem signature or the word `by`, and do " + "not use `sorry`/`admit`: the goal must be fully closed." + } + ) + + def __post_init__(self) -> None: + if _SORRY.search(self.tactics): + raise ValueError( + "the proof uses `sorry`/`admit`; the goal must be closed for real " + "(sorry is only allowed for the separately-proposed lemmas)" + ) + ctx = LEAN_CTX.get() + result = ctx.kernel.compile(assemble(ctx.decl, self.tactics, ctx.preamble)) + if not result.ok: + raise ValueError( + "Lean rejected this proof. Fix it against the compiler output below " + f"(you may call `check` to iterate):\n{result.messages}" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class ProposedLemma: + """One intermediate lemma a blueprint proposes: a Lean declaration header the + sketch may assume. ``__post_init__`` certifies the *statement* type-checks (as a + ``sorry`` stub) so a malformed or ill-typed lemma is fed back before it becomes a + subgoal -- the statement must at least be a well-formed proposition, even though + its proof is deferred.""" + + name: str = dataclasses.field( + metadata={ + "description": "A fresh Lean identifier for the lemma (snake_case, unique " + "within this decomposition), referenced by name from the sketch." + } + ) + decl: str = dataclasses.field( + metadata={ + "description": "The lemma's Lean declaration header WITHOUT the name or " + "`:= ...`, i.e. the binders and proposition: e.g. `(n : ℕ) : 0 < n + 1`. " + "It must type-check as a standalone statement." + } + ) + rationale: str = dataclasses.field( + metadata={ + "description": "Why proving this lemma helps -- what it lets the sketch do, " + "and why it is strictly simpler / more general than the goal." + } + ) + + def header(self) -> str: + """The full stub header ``theorem `` for the compile preamble.""" + return f"theorem {self.name} {self.decl}" + + def __post_init__(self) -> None: + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_']*", self.name): + raise ValueError(f"lemma name {self.name!r} is not a valid Lean identifier") + ctx = LEAN_CTX.get() + # The statement must type-check; its proof may be deferred (`sorry`). + stub = f"{ctx.preamble}\n\n{self.header()} := sorry\n" + result = ctx.kernel.compile(stub) + if not result.ok: + raise ValueError( + f"the lemma statement `{self.name} {self.decl}` does not type-check. " + f"Fix the statement against Lean's output:\n{result.messages}" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Blueprint: + """The informal decomposition: a natural-language plan plus the intermediate + lemmas it proposes. The sketch (a ``LeanProof``) is the formal counterpart that + proves the goal assuming these lemmas.""" + + plan: str = dataclasses.field( + metadata={ + "description": "The informal proof blueprint in a few sentences: how the " + "goal reduces to the proposed lemmas." + } + ) + lemmas: list[ProposedLemma] + + def __post_init__(self) -> None: + if not self.lemmas: + raise ValueError( + "a decomposition must propose at least one lemma; if the goal needs " + "no lemmas it should be proved directly, not decomposed" + ) + names = [lm.name for lm in self.lemmas] + if len(set(names)) != len(names): + raise ValueError(f"proposed lemma names are not unique: {names}") + + +@pydantic.dataclasses.dataclass(frozen=True) +class ReviewVerdict: + """The decomposition reviewer's judgment -- the paper's planning-level search + filter (Sec. 2.5 / Figure 3).""" + + simplifies: bool = dataclasses.field( + metadata={ + "description": "True only if every proposed lemma is genuinely simpler or " + "more general than the goal and plausibly provable; false if any lemma " + "merely restates the goal or is no easier than it." + } + ) + reason: str + + +# --------------------------------------------------------------------------- +# The compiler-in-the-loop base class. The `check` tool is defined here, so it +# reaches the FormalProver and SketchAgent (via the Agent MRO) and is invisible to +# the closed-book NLProver, BlueprintAgent, and Reviewer. +# --------------------------------------------------------------------------- + + +class LeanAgent(Agent): + """Base for agents that write Lean against the live compiler. The ``check`` tool + defined here compiles a candidate tactic block for the goal in scope and returns + Lean's messages, so a formalizing agent can iterate against real compiler feedback + before committing an answer -- the paper's continuous compiler interaction. Agents + that only plan or review do not subclass ``LeanAgent``, so the tool never enters + their lexical scope.""" + + @Tool.define + def check(self, tactics: str) -> str: + """Compile `` := by `` with Lean and return the + compiler's output: an empty/clean result means it is accepted; otherwise the + errors and the remaining goal state. Use this to test tactics and read the + goal before you commit a final proof. (You may also use Mathlib's own + ``exact?`` / ``apply?`` inside ``tactics`` to search for premises.)""" + ctx = LEAN_CTX.get() + if _SORRY.search(tactics): + return ( + "Refused: `tactics` contains `sorry`/`admit`; the goal must be closed." + ) + result = ctx.kernel.compile(assemble(ctx.decl, tactics, ctx.preamble)) + if result.ok: + return "Lean accepts this proof (no errors)." + return f"Lean output:\n{result.messages}" + + +# --------------------------------------------------------------------------- +# Stage 1 -- direct formalization: NL prover (informal) -> formal prover (Lean). +# --------------------------------------------------------------------------- + + +class NLProver(Agent): + """You are the informal reasoner. You write a short, rigorous natural-language + proof of a statement -- the mathematical argument, not Lean code -- for a + formalizer to translate. Closed-book: you hold no compiler tool.""" + + @Template.define + def argue(self, goal: str, context: str) -> str: + """Give a concise but rigorous informal proof of the following statement. + State the key steps a formal proof would need (case splits, inductions, + lemmas invoked). Do not write Lean. + + Statement: + {goal} + + Context that may help (available lemmas already proved, and the shape of the + problem): + {context} + """ + + +class FormalProver(LeanAgent): + """You are the formal prover. You translate an informal argument into a Lean 4 + tactic proof and make it compile, using the ``check`` tool to iterate against the + real compiler. You prefer short, robust proofs (``simp``, ``omega``, ``ring``, + ``induction``, ``exact?``) and you never leave a ``sorry``.""" + + @Template.define + def formalize(self, goal: str, informal: str, context: str) -> LeanProof: + """Prove the goal below in Lean 4 with Mathlib by returning the tactic block + (what follows ``:= by``). Translate the informal argument, then use ``check`` + to compile and fix it against Lean's output until it is accepted. The proof + must fully close the goal -- no ``sorry``. + + Goal declaration (your tactics complete `` := by ...``): + {goal} + + Informal argument to formalize: + {informal} + + Context (lemmas already proved and in scope; you may cite them by name): + {context} + """ + + +# --------------------------------------------------------------------------- +# Stage 2 -- decomposition: blueprint (informal) -> reviewer -> sketch (Lean). +# --------------------------------------------------------------------------- + + +class BlueprintAgent(Agent): + """You are the blueprint planner. When a goal resists direct proof, you propose + a decomposition: intermediate lemmas that are each strictly simpler or more + general than the goal, such that the goal follows easily once they hold. Closed- + book: you plan in mathematics, not against the compiler.""" + + @Template.define + def draft(self, goal: str, context: str, feedback: str) -> Blueprint: + """The goal below could not be proved directly within budget. Draft a proof + blueprint: an informal plan plus a small set of intermediate lemmas that make + the goal easy to prove. Each lemma must be genuinely simpler or more general + than the goal -- never a restatement of it -- and should be broadly useful. + Give each lemma a fresh Lean identifier and a well-formed statement. + + Goal declaration: + {goal} + + Context (lemmas already proved and in scope -- prefer reusing these to + proposing new ones): + {context} + + Feedback from prior attempts (empty on the first try): + {feedback} + """ + + +class SketchAgent(LeanAgent): + """You are the sketch formalizer. Given a blueprint, you write a Lean tactic + proof of the goal that *assumes the proposed lemmas* (they are in scope as + hypotheses you may cite by name). Your tactics themselves must be ``sorry``-free: + the goal must reduce to the lemmas. Use ``check`` to compile against the real + Lean, where the proposed lemmas are present as stubs.""" + + @Template.define + def sketch(self, goal: str, blueprint: str, context: str) -> LeanProof: + """Prove the goal below assuming the blueprint's lemmas. Return the tactic + block (what follows ``:= by``); you may reference each proposed lemma by its + name as an already-proved fact. Your tactics must not use ``sorry`` -- only the + lemmas are deferred. Use ``check`` to compile and fix against Lean's output. + + Goal declaration: + {goal} + + Blueprint (plan and the lemmas now in scope, by name): + {blueprint} + + Context (other lemmas already proved and in scope): + {context} + """ + + +class Reviewer(Agent): + """You are the decomposition reviewer -- a planning-level search filter. Compiler + verification only checks that a sketch is well-typed, not that its decomposition + makes progress: a sketch can compile while proposing a subgoal no simpler than the + goal (e.g. one syntactically equivalent to it). You reject such non-simplifying + decompositions so search does not waste effort on them.""" + + @Template.define + def review(self, goal: str, blueprint: str) -> ReviewVerdict: + """Judge whether this decomposition genuinely simplifies proving the goal. + Reject it if any proposed lemma merely restates the goal, is no easier than + it, or does not plausibly advance the proof. Accept only a decomposition whose + lemmas are each strictly simpler or more general than the goal and together + make it easy. + + Goal declaration: + {goal} + + Proposed decomposition: + {blueprint} + """ + + +# --------------------------------------------------------------------------- +# The AND-OR DAG -- proof progress and hierarchical memoization (paper Sec. 2.3). +# --------------------------------------------------------------------------- + + +def _norm(text: str) -> str: + """Collapse whitespace so two statements that differ only in spacing share a + memoization key. (Textual, not Lean-defeq: enough for the reuse story.)""" + return " ".join(text.split()) + + +@dataclasses.dataclass +class GoalNode: + """An OR node: a goal (or lemma) to prove. ``decl`` is its Lean header + ``theorem ``; once ``proved``, ``tactics`` is the accepted tactic + block, and the node is a reusable Lean declaration in every downstream preamble.""" + + name: str + decl: str # "theorem : " + proved: bool = False + attempted: bool = False + tactics: str | None = None + reused: bool = False + + def declaration(self) -> str: + """The full proved Lean declaration, for the reuse preamble and final file.""" + assert self.proved and self.tactics is not None + body = textwrap.indent(self.tactics.strip(), " ") + return f"{self.decl} := by\n{body}" + + +@dataclasses.dataclass +class ProofDAG: + """The proof graph: OR nodes keyed by normalized statement (memoization), plus a + monotonically-growing preamble of proved lemma declarations that every subsequent + compile reuses. The ``state_reader``/``state_writer`` of the paper are this + object's read/commit methods.""" + + kernel: LeanKernel + nodes: dict[str, GoalNode] = dataclasses.field(default_factory=dict) + # Proved nodes in completion order. A node is proved only after its children, so + # this order is topological (dependencies first) -- the order the reuse preamble + # and the final assembly must emit declarations in. + proof_order: list[GoalNode] = dataclasses.field(default_factory=list) + _counter: int = 0 + + def fresh_name(self, hint: str) -> str: + self._counter += 1 + slug = re.sub(r"[^A-Za-z0-9_]", "_", hint).strip("_")[:24] or "lemma" + return f"leap_{self._counter}_{slug}" + + def get_or_add(self, sig: str, name_hint: str) -> tuple[GoalNode, bool]: + """Look a goal up by normalized signature; create its OR node if new. Returns + (node, is_new). A hit on a proved node is the memoization payoff.""" + key = _norm(sig) + if key in self.nodes: + return self.nodes[key], False + name = self.fresh_name(name_hint) + node = GoalNode(name=name, decl=f"theorem {name} {sig}") + self.nodes[key] = node + return node, True + + def mark_proved(self, node: GoalNode, tactics: str) -> None: + """Commit a node's accepted proof and record it in topological order.""" + node.proved, node.tactics = True, tactics + self.proof_order.append(node) + + def proved_preamble(self) -> str: + """PRELUDE + every proved lemma's real declaration in dependency order, so any + compile reuses the whole proved library (real Lean-level lemma sharing).""" + decls = [n.declaration() for n in self.proof_order] + return PRELUDE + ("\n\n".join(decls) + "\n\n" if decls else "") + + def context_digest(self) -> str: + """A short human/model-readable list of proved lemmas in scope (state_reader).""" + if not self.proof_order: + return "(no lemmas proved yet)" + return "\n".join(f"- {n.decl}" for n in self.proof_order) + + +# --------------------------------------------------------------------------- +# Verification-guided proof search: direct proof, else decompose. DFS + backtrack. +# --------------------------------------------------------------------------- + + +def _sig_of_lemma(lm: ProposedLemma) -> str: + """A proposed lemma's signature (binders : prop) -- its memoization key.""" + return lm.decl.strip() + + +def try_direct(dag: ProofDAG, node: GoalNode, sig: str, depth: int) -> bool: + """Attempt a direct proof: informal argument -> Lean formalization, certified by + the compiler on decode. Returns True and records the tactics on success; on + failure (retries exhausted without a compiling proof) returns False so the caller + decomposes. This is the paper's direct-formalization path with REVISER feedback + (here, ``RetryLLMHandler``).""" + ind = " " * depth + ctx = LeanContext(dag.kernel, dag.proved_preamble(), node.decl) + informal = with_ctx(ctx, lambda: NLProver().argue(node.decl, dag.context_digest())) + try: + proof = with_ctx( + ctx, + lambda: FormalProver().formalize(node.decl, informal, dag.context_digest()), + ) + except Exception as exc: # retries exhausted without a compiling proof + print(f"{ind}[direct] no compiling proof ({type(exc).__name__}); decomposing") + return False + dag.mark_proved(node, proof.tactics) + print( + f"{ind}[direct] proved `{node.name}` ({len(proof.tactics.splitlines())} tactic lines)" + ) + return True + + +def decompose( + dag: ProofDAG, node: GoalNode, sig: str, ancestors: frozenset[str], depth: int +) -> bool: + """Blueprint -> review -> sketch -> recurse. A decomposition is committed only if + the reviewer finds it simplifying, the state_writer finds it acyclic, the sketch + compiles (assuming the lemmas), and every child subgoal is then proved. On any + failure it backtracks and re-drafts, up to a bound (DFS with backtracking).""" + ind = " " * depth + feedback = "" + for attempt in range(1, MAX_BLUEPRINTS + 1): + # Blueprint (informal plan + proposed lemma statements). ProposedLemma decode + # type-checks each statement against the proved preamble, so it needs the ctx. + base = LeanContext(dag.kernel, dag.proved_preamble(), node.decl) + try: + blueprint = with_ctx( + base, + lambda: BlueprintAgent().draft( + node.decl, dag.context_digest(), feedback + ), + ) + except Exception as exc: + print(f"{ind}[blueprint {attempt}] draft failed ({type(exc).__name__})") + continue + bp_text = _render_blueprint(blueprint) + + # Reviewer: reject a decomposition that does not simplify (Figure 3). + verdict = Reviewer().review(node.decl, bp_text) + if not verdict.simplifies: + print(f"{ind}[review {attempt}] rejected: {verdict.reason}") + feedback = ( + f"A reviewer rejected the previous decomposition: {verdict.reason}" + ) + continue + + # state_writer: reject any subgoal that would reintroduce an ancestor -- keep + # the DAG acyclic (also catches the Figure-3 "subgoal == parent" pathology). + cyclic = [ + lm.name for lm in blueprint.lemmas if _norm(_sig_of_lemma(lm)) in ancestors + ] + if cyclic: + print(f"{ind}[state_writer {attempt}] rejected cyclic subgoals {cyclic}") + feedback = ( + f"These proposed lemmas restate an ancestor goal (a cycle): {cyclic}. " + "Propose strictly simpler, non-circular lemmas." + ) + continue + + # Bind each proposed lemma to a DAG node (memoization-aware): a lemma whose + # statement is already a node reuses that node -- and its name -- so the sketch, + # the stubs, and the eventually-stored proof all agree on one identifier. This + # is what makes the assembled proof reference real, in-scope declarations. + lemma_nodes = [ + dag.get_or_add(_sig_of_lemma(lm), lm.name)[0] for lm in blueprint.lemmas + ] + + # Sketch: prove the goal assuming the lemmas. Already-proved lemmas are in the + # proved preamble; the rest are installed as `sorry` stubs under their DAG names. + stubs = "\n\n".join( + f"{nd.decl} := sorry" for nd in lemma_nodes if not nd.proved + ) + sketch_ctx = LeanContext( + dag.kernel, + dag.proved_preamble() + (stubs + "\n\n" if stubs else ""), + node.decl, + ) + sketch_bp = _render_blueprint(blueprint, lemma_nodes) + try: + sketch = with_ctx( + sketch_ctx, + lambda: SketchAgent().sketch( + node.decl, sketch_bp, dag.context_digest() + ), + ) + except Exception as exc: + print(f"{ind}[sketch {attempt}] no compiling sketch ({type(exc).__name__})") + feedback = "The sketch did not compile even assuming the lemmas; simplify the plan." + continue + print( + f"{ind}[sketch {attempt}] compiles assuming {len(blueprint.lemmas)} lemma(s)" + ) + + # Recurse on each subgoal, sharing the DAG (memoization across branches). + child_ancestors = ancestors | {_norm(sig)} + all_proved = True + for lm in blueprint.lemmas: + if not prove(dag, _sig_of_lemma(lm), lm.name, child_ancestors, depth + 1): + all_proved = False + break + if not all_proved: + print(f"{ind}[decompose {attempt}] a subgoal failed; backtracking") + feedback = "A proposed lemma could not be proved; propose different lemmas." + continue + + # All children proved -> the AND node succeeds -> the parent is proved. Its + # reusable proof is the sketch over the now-real (not sorry) lemma preamble. + dag.mark_proved(node, sketch.tactics) + print(f"{ind}[decompose {attempt}] all subgoals proved -> `{node.name}` proved") + return True + + print(f"{ind}[decompose] exhausted {MAX_BLUEPRINTS} blueprints for `{node.name}`") + return False + + +def prove( + dag: ProofDAG, sig: str, name_hint: str, ancestors: frozenset[str], depth: int +) -> bool: + """Prove a goal (statement ``sig`` = ``binders : prop``): memo hit, else direct, + else decompose. Shared ``dag`` gives hierarchical memoization; ``ancestors`` + enforces acyclicity.""" + ind = " " * depth + node, _ = dag.get_or_add(sig, name_hint) + if node.proved: # memoization hit: a lemma already proved in another branch + node.reused = True + print(f"{ind}[memo] reuse proved lemma `{node.name}`: {_norm(sig)[:70]}") + return True + if _norm(sig) in ancestors: # this goal is its own ancestor -> a cycle + print(f"{ind}[cycle] `{node.name}` restates an ancestor; abandoning branch") + return False + if node.attempted: # tried before and not proved; don't loop on it again + print(f"{ind}[skip] `{node.name}` was already attempted and failed") + return False + node.attempted = True + print(f"{ind}[goal] {node.name}: {_norm(sig)[:80]}") + + # Normally: direct first, decompose on failure. ``--decompose-root`` forces the + # root to decompose so the blueprint/reviewer/sketch/memoization path is exercised + # even when a strong model could one-shot it (a labeled demo, not the paper's flow). + force = FORCE_DECOMPOSE_ROOT and depth == 0 + if not force and try_direct(dag, node, sig, depth): + return True + if depth >= MAX_DEPTH: + print(f"{ind}[depth] max decomposition depth reached for `{node.name}`") + return False + return decompose(dag, node, sig, ancestors, depth) + + +def _render_blueprint(bp: Blueprint, nodes: list[GoalNode] | None = None) -> str: + """Render a blueprint for a prompt. When ``nodes`` is given (one per lemma, in + order), lemmas are named by their DAG identifier -- the name the sketch must cite + and under which the proof is stored -- so the sketch references real declarations.""" + names = ( + [nd.name for nd in nodes] + if nodes is not None + else [lm.name for lm in bp.lemmas] + ) + lines = [bp.plan, "", "Proposed lemmas (in scope, cite by the name shown):"] + for name, lm in zip(names, bp.lemmas): + lines.append(f"- {name} : {lm.decl} -- {lm.rationale}") + return "\n".join(lines) + + +# Search bounds. +MAX_BLUEPRINTS = 3 # decomposition re-drafts before a node is abandoned (backtracking) +MAX_DEPTH = 3 # deepest decomposition nesting +FORCE_DECOMPOSE_ROOT = False # set by --decompose-root: skip the root's direct attempt + + +# --------------------------------------------------------------------------- +# The pipeline +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class LeapResult: + proved: bool + dag: ProofDAG + root: GoalNode + + +def run_leap(root_sig: str, kernel: LeanKernel) -> LeapResult: + """Register the root theorem and drive the DFS: direct-then-decompose over the + shared AND-OR DAG. Returns the DAG so the caller can inspect memoization and emit + the assembled proof.""" + dag = ProofDAG(kernel=kernel) + proved = prove(dag, root_sig, "root", frozenset(), 0) + root = dag.nodes[_norm(root_sig)] + return LeapResult(proved=proved, dag=dag, root=root) + + +def assemble_full_proof(dag: ProofDAG) -> str: + """Emit the whole proof tree as one Lean file: PRELUDE + every proved lemma + + the root, ordered so dependencies precede uses (proved-order suffices since a + lemma is proved before the parent that uses it). Compiling this with no ``sorry`` + is the end-to-end check -- like ``scientist_one``'s audit re-deriving its + evidence.""" + parts = [PRELUDE.strip(), ""] + for n in dag.proof_order: # topological: dependencies precede uses; root is last + parts += [n.declaration(), ""] + return "\n".join(parts) + + +# --------------------------------------------------------------------------- +# Targets: small theorems whose proof benefits from a lemma decomposition. The +# statement is the Lean signature `(binders) : proposition`; LEAP supplies the name. +# --------------------------------------------------------------------------- + +TARGETS: dict[str, str] = { + # Sum of the first n odd numbers is n^2. Direct: induction + `Finset.sum_range_succ` + # + `ring`. A natural decomposition proves the successor step as its own lemma. + "odd_sum": r"(n : ℕ) : (∑ i ∈ Finset.range n, (2 * i + 1)) = n ^ 2", + # Gauss sum, doubled to stay in ℕ. Induction; the step is a clean sub-lemma. + "gauss": r"(n : ℕ) : (2 * ∑ i ∈ Finset.range (n + 1), i) = n * (n + 1)", + # A divisibility fact that invites a two-lemma decomposition (parity of n*(n+1) + # feeding 6 ∣ n*(n+1)*(n+2)); harder, exercises deeper decomposition. + "div6": r"(n : ℕ) : 6 ∣ n * (n + 1) * (n + 2)", +} + +# A trivial theorem used to validate the toolchain without any LLM. +SANITY = r"(n : ℕ) : n + 0 = n" + + +def check_toolchain(kernel: LeanKernel) -> None: + """Compile a trivial theorem (no LLM) to confirm Lean+Mathlib is wired up.""" + if not kernel.available(): + print( + f"Lean project not found/built at {kernel.project!r}. Build it once:\n" + " elan default stable # if elan is installed\n" + f" cd {kernel.project} && lake exe cache get && lake build" + ) + return + print(f"Compiling a trivial theorem via {kernel.project} ...") + src = f"{PRELUDE}\ntheorem leap_sanity {SANITY} := by simp\n" + result = kernel.compile(src) + print("Toolchain OK." if result.ok else f"Toolchain FAILED:\n{result.messages}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--target", + choices=list(TARGETS), + default="odd_sum", + help="Which theorem to prove", + ) + parser.add_argument( + "--statement", + type=str, + default=None, + help="A custom Lean signature `(binders) : prop` to prove instead of --target", + ) + parser.add_argument( + "--project", + type=str, + default=LEAN_PROJECT, + help="Path to the Lean+Mathlib lake project", + ) + parser.add_argument( + "--timeout", + type=float, + default=120.0, + help="Per-compile timeout in seconds", + ) + parser.add_argument( + "--check-toolchain", + action="store_true", + help="Skip the pipeline; compile a trivial theorem to validate Lean+Mathlib", + ) + parser.add_argument( + "--decompose-root", + action="store_true", + help="Force the root goal to decompose (skip its direct proof), to demonstrate " + "the blueprint/reviewer/sketch/memoization path even on an easy target", + ) + args = parser.parse_args() + + global FORCE_DECOMPOSE_ROOT + FORCE_DECOMPOSE_ROOT = args.decompose_root + kernel = LeanKernel(project=args.project, timeout=args.timeout) + + if args.check_toolchain: + check_toolchain(kernel) + return + + if not kernel.available(): + raise SystemExit( + f"Lean project not built at {args.project!r}; run with --check-toolchain " + "for build instructions." + ) + + sig = args.statement or TARGETS[args.target] + print(f"Proving: {sig}\n") + + result = run_leap(sig, kernel) + + print("\n" + "=" * 72) + proved = [n for n in result.dag.nodes.values() if n.proved] + reused = [n for n in result.dag.nodes.values() if n.reused] + print( + f"DAG: {len(result.dag.nodes)} goal node(s), {len(proved)} proved, " + f"{len(reused)} reused via memoization." + ) + for n in result.dag.nodes.values(): + mark = "proved" if n.proved else "OPEN" + extra = " (reused)" if n.reused else "" + print(f" [{mark}]{extra} {n.decl}") + + if not result.proved: + raise SystemExit( + "\nLEAP did not close the root goal (as the paper notes, " + "one-shot formal proving is hard; try another --target)." + ) + + # End-to-end verification: compile the whole assembled proof tree with no sorry. + full = assemble_full_proof(result.dag) + print("\nAssembled proof; recompiling the whole tree end-to-end ...") + final = kernel.compile(full) + if final.ok: + print("VERIFIED: the complete proof compiles under Lean with no `sorry`.\n") + print(full) + else: + raise SystemExit(f"Assembled proof failed to recompile:\n{final.messages}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoresearch/__init__.py b/docs/source/llm_examples/autoresearch/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/autoresearch/illustration.py b/docs/source/llm_examples/autoresearch/illustration.py new file mode 100644 index 000000000..88ec51114 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/illustration.py @@ -0,0 +1,689 @@ +"""PaperBanana: reference-driven academic illustration as a 5-agent pipeline. + +Implements the core of "PaperBanana: Automating Academic Illustration for AI +Scientists" (arXiv:2601.23265). The paper maps a *source context* S (a methodology +or plot description, including the data) and a *communicative intent* C (a figure +caption) to an illustration ``I = f(S, C, E)``, optionally guided by a reference +set E, via five specialized agents in two phases: a *Linear Planning Phase* +(Retriever -> Planner -> Stylist) that synthesizes a stylistically optimized +description P*, and an *Iterative Refinement Loop* (T=3) in which a Visualizer +renders P and a Critic inspects the render and refines it. + +We implement PaperBanana's 5-agent architecture on its **code-based +statistical-plot path** (paper Sec. 5.5), which makes the Visualizer<->Critic loop +*real*: the Visualizer writes executable Matplotlib code, ``matplotlib`` renders it +to a PNG, and a vision model critiques that actual PNG against S and C. The raster +methodology-diagram path needs an image-generation model (Nano-Banana-Pro / +GPT-Image) and is out of scope. Each agent falls out of an ordinary effectful idiom: + + * Retriever -- generative retrieval as decode-time certification. A ``Retrieval`` + names keys of exemplars from the fixed reference set R; ``__post_init__`` rejects + any key that does not resolve in the immutable ``REFERENCES`` constant, so a + hallucinated selection is fed back by ``RetryLLMHandler`` (the certification + idiom of ``scholar_peer.py``/``scientist_one.py``). Because R is a module + constant, not per-run mutable state, the check reads it directly -- no ContextVar + needed. A retrieval ``Tool`` surfaces the candidate metadata as a strong + ``list[PlotExemplar]``, the way ``scientist_one``'s tools return domain types. + + * Planner -- in-context learning, toolless. It reads the retrieved exemplars and + transcribes S's data table into a structured, data-bearing ``PlotDescription``, + so the numbers the Visualizer draws are carried explicitly and stay checkable. + + * Stylist -- a plan threaded through the pipeline. It synthesizes an + ``AestheticGuideline`` G from R, then restyles P into P* (a ``StyledPlot`` + bundling the data-bearing description with G), exactly the typed-plan-as- + orchestration shape of ``paper_orchestra.py``'s Outline. + + * Visualizer -- code synthesis that must actually render. Its Template returns a + ``Plot`` (a ``Callable`` the harness compiles, as in ``scientist_one``'s + ``Solver -> Solution``): a pure nullary closure that builds and returns a fresh + ``Figure`` via matplotlib's object-oriented API. A render-doctest in its docstring + calls the synthesized ``plot()``, so a plot whose code raises when it runs fails + the doctest and is fed back by ``RetryLLMHandler`` -- "the plot must render" is + grounding by construction, and the doctest runs the code on a *different* plan, + forcing it to read its data from the closed-over plan rather than + hardcode. + + * Critic -- the loop made real, multimodally. It receives the rendered + ``PIL.Image.Image`` (the image-input idiom of ``image_input.py``), inspects it + against S and C for factual misalignments and visual glitches, and returns a + refined ``StyledPlot`` plus the concrete issues it saw. The Visualizer<->Critic + loop is a plain Python ``for`` loop over these two Templates. + + * Judge -- referenced, multimodal, hierarchical. It compares two rendered PNGs + (round-0 P* vs the final round-T render -- the paper's Critic-on/off ablation) + against S and C on four dimensions (Faithfulness, Conciseness, Readability, + Aesthetics), each win/tie/loss, and plain Python aggregates them under the + paper's hierarchical rule: faithfulness+readability are primary and decide the + winner; conciseness+aesthetics break primary ties ("show the truth"). + +Demonstrates: +- Code synthesis whose return ``Callable`` must render: the model writes a pure + nullary closure that builds and returns a ``Figure`` via matplotlib's OO API, and a + doctest turns "the plot actually renders" into a decode-time contract fed back by + ``RetryLLMHandler`` +- A real multimodal refinement loop: matplotlib renders a PNG a vision model + critiques, then the plan is regenerated -- the Visualizer<->Critic loop, not simulated +- Decode-time certification of a retrieval selection against an immutable reference + set, read directly (no ContextVar) because the set is a module constant +- A typed plan (``StyledPlot``) threaded through the pipeline as orchestration data +- A referenced multimodal LLM judge with hierarchical win/tie/loss aggregation, + stateless (a fresh instance per call) +- Per-field guidance carried on the types via ``field(metadata={"description": ...})`` +""" + +# Simplifications vs. the source: +# - The raster methodology-diagram path -- PaperBanana's headline -- is out of scope: +# it needs an image-generation model (Nano-Banana-Pro / GPT-Image). We implement the +# paper's own code-based statistical-plot path (Sec. 5.5), where the Visualizer emits +# Matplotlib and the loop is a real render+critique cycle rather than image gen. +# - Static reference corpus, not live/web-scale retrieval. ``REFERENCES`` is a tiny +# in-memory set of *textual* structure/style descriptors (no exemplar images), so +# the Retriever ranks over metadata; this shows the pipeline's shape, not retrieval +# at scale, and the "prioritize visual structure over topic" instruction is only +# gestured at without real reference images. +# - One task, not PaperBananaBench. The paper curates 292 evaluation cases; here a +# single planted illustration task runs end to end, as the sibling examples do. +# - No human-reference comparison. The paper's VLM-as-a-Judge scores against a +# human-drawn figure; lacking one, we use the paper's own Critic-on/off ablation as +# the referenced pair (round-0 P* vs final round-T), which isolates the loop's value. + +import argparse +import collections.abc +import dataclasses +import pathlib +import tempfile +import typing + +import pydantic +from matplotlib.figure import Figure +from PIL import Image + +from effectful.handlers.llm import Agent, Template, Tool + +# A field's ``metadata={"description": ...}`` is inlined by pydantic into that +# field's JSON schema, which the harness renders into the system prompt as part of +# a template's argument (and structured-output) spec. So per-field guidance reaches +# the model *through the type* -- used below only where the field name and type +# don't already say it, so no prompt has to repeat it. + +type ChartType = typing.Literal[ + "grouped_bar", "stacked_bar", "line", "scatter", "heatmap" +] + +# A plotting function the Visualizer writes: a pure nullary closure that builds and +# returns a fresh matplotlib ``Figure`` via the object-oriented API (no pyplot, no +# global figure registry, no side effects). The harness compiles the model's code +# into one of these; the closure captures the round's ``plan``. +type PlottingFn = collections.abc.Callable[[], Figure] + + +# --------------------------------------------------------------------------- +# The reference set R -- the fixed corpus of exemplars the Retriever ranks over. +# In PaperBanana each exemplar is a triplet (S, C, I) with a real reference image; +# here it is textual structure/style metadata (no images), so an exemplar has a +# stable key a selection can be certified against. R is an immutable module +# constant, so the certification reads it directly (contrast ``scientist_one``, +# whose per-run mutable Workspace needs a ContextVar). +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class PlotExemplar: + """One reference exemplar: its chart type, research domain, and -- kept separate + so the Retriever can weight them as the paper does -- its *visual structure* + (prioritized) versus its *topic* and aesthetic style.""" + + key: str + chart_type: ChartType + domain: str + caption: str + structure_notes: str = dataclasses.field( + metadata={ + "description": "The visual/structural composition (axes, grouping, marks, " + "legend, layout) independent of subject matter -- what the Retriever " + "weights above topic when matching." + } + ) + style_notes: str + + +REFERENCES: dict[str, PlotExemplar] = { + e.key: e + for e in [ + PlotExemplar( + key="grouped_bar_benchmark", + chart_type="grouped_bar", + domain="ML benchmarking", + caption="Accuracy of several methods across benchmarks.", + structure_notes="Bars clustered by benchmark on the x-axis, one colored " + "bar per method within each cluster; shared y-axis starting at 0; a legend " + "keys color to method.", + style_notes="Categorical palette, one hue per method; light horizontal " + "gridlines; top and right spines removed; value labels above bars.", + ), + PlotExemplar( + key="line_scaling", + chart_type="line", + domain="scaling laws", + caption="A metric as a function of training scale.", + structure_notes="Several monotone lines share x (steps/size, often log) " + "and y (the metric); one line per method, markers at measured points.", + style_notes="Distinct hue+marker per line; faint grid; legend inside the " + "plot; no chartjunk.", + ), + PlotExemplar( + key="scatter_tradeoff", + chart_type="scatter", + domain="efficiency analysis", + caption="Accuracy vs. cost trade-off across methods.", + structure_notes="Points in an x=cost / y=quality plane, one marker per " + "method; a Pareto frontier implied toward the upper-left.", + style_notes="One hue per method, labeled points; equal-weight axes; " + "minimal grid.", + ), + PlotExemplar( + key="heatmap_ablation", + chart_type="heatmap", + domain="ablation study", + caption="A metric over a grid of two design choices.", + structure_notes="A matrix of cells indexed by two categorical axes, cell " + "color encoding the metric; a colorbar legend; cells annotated with values.", + style_notes="Sequential colormap; annotated cells; square aspect.", + ), + PlotExemplar( + key="stacked_bar_composition", + chart_type="stacked_bar", + domain="component analysis", + caption="Contribution of components to a total per setting.", + structure_notes="One bar per setting on x, segments stacked to a total on " + "y, each segment a component; a legend keys color to component.", + style_notes="Sequential/categorical stack palette; legend outside; totals " + "labeled atop each bar.", + ), + ] +} + + +# --------------------------------------------------------------------------- +# Inputs: the task the illustration must satisfy. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class IllustrationTask: + """The task ``(S, C)``: a source context and a communicative intent. ``S`` + embeds the actual data (as a small text table) so faithfulness is checkable; + ``C`` is the figure caption that fixes the illustration's scope and focus.""" + + source_context: str + intent: str + + +# --------------------------------------------------------------------------- +# Structured artifacts crossing between agents. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class Retrieval: + """The Retriever's selection E: keys of the exemplars from R that best match the + task by diagram type and domain (visual structure weighted over topic).""" + + selected: list[str] = dataclasses.field( + metadata={ + "description": "Keys of the chosen exemplars; each MUST resolve in the " + "reference set R (certified at decode time), so a hallucinated key is " + "rejected and fed back." + } + ) + rationale: str + + def __post_init__(self) -> None: + unknown = [k for k in self.selected if k not in REFERENCES] + if unknown: + raise ValueError( + f"retrieval selected unknown exemplar keys {unknown}; choose only " + f"keys that exist in the reference set (available: {sorted(REFERENCES)})" + ) + if not self.selected: + raise ValueError("select at least one exemplar from the reference set") + + +@pydantic.dataclasses.dataclass(frozen=True) +class DataSeries: + """One data series (a method / line / stack): a name and its numeric values.""" + + name: str + values: list[float] = dataclasses.field( + metadata={ + "description": "One value per category, in the SAME order as the " + "description's ``categories`` -- the transcribed numbers from S the plot " + "must reproduce exactly." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class PlotDescription: + """The Planner's description P of the target plot: its type, labels, and -- carried + explicitly so the Visualizer's code stays faithful -- the exact data from S.""" + + chart_type: ChartType + title: str + x_label: str + y_label: str + categories: list[str] = dataclasses.field( + metadata={ + "description": "The groups along the x-axis (e.g. datasets); the " + "Visualizer's code iterates these and each series aligns to them by order." + } + ) + series: list[DataSeries] + notes: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class AestheticGuideline: + """The Stylist's synthesized guideline G, one directive per aesthetic dimension + (the paper's palette / shapes / lines / layout / typography / icons), read off R + and specialized to statistical plots.""" + + palette: list[str] = dataclasses.field( + metadata={ + "description": "Ordered color specs (hex like '#4C72B0' or matplotlib " + "names), one per series, applied in order." + } + ) + marks_and_containers: str + lines_and_arrows: str + layout: str + typography: str + icons: str = dataclasses.field( + metadata={ + "description": "Any small glyphs/markers or annotation style; 'none' for a " + "plain statistical plot." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class StyledPlot: + """P* -- the plan the Visualizer renders and the Critic refines, threaded through + the refinement loop. Bundles the data-bearing description with the aesthetic + guideline G and the concrete directives that restyle P into P*.""" + + description: PlotDescription + guideline: AestheticGuideline + directives: str = dataclasses.field( + metadata={ + "description": "Concrete restyling instructions applying G to this " + "description -- what colors/spines/gridlines/labels the Visualizer should use." + } + ) + + def __str__(self) -> str: + """Render the plan to a compact, exact brief -- data first, so the Visualizer + (and a human) sees the precise numbers it must draw.""" + d = self.description + lines = [ + f"{d.chart_type} titled {d.title!r}", + f" x-axis ({d.x_label}): {d.categories}", + f" y-axis: {d.y_label}", + " series:", + ] + lines += [f" - {s.name}: {s.values}" for s in d.series] + g = self.guideline + lines += [ + f" planner notes: {d.notes}", + f" palette: {g.palette}", + f" marks/containers: {g.marks_and_containers}", + f" lines/arrows: {g.lines_and_arrows}", + f" layout: {g.layout}", + f" typography: {g.typography}", + f" icons: {g.icons}", + f" style directives: {self.directives}", + ] + return "\n".join(lines) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Critique: + """The Critic's verdict on one rendered plot: the concrete problems it saw and a + refined plan addressing them.""" + + issues: list[str] = dataclasses.field( + metadata={ + "description": "What the Critic targets in the RENDERED plot, per the " + "paper: factual misalignments (wrong/missing numbers or labels vs. S and " + "C), visual glitches, OR areas for improvement (readability/aesthetics). " + "Leave empty ONLY if the plot is already publication-ready with nothing to " + "improve." + } + ) + refined: StyledPlot + + +# --------------------------------------------------------------------------- +# The Visualizer's render path: compile the model's code, draw it to a real PNG. +# --------------------------------------------------------------------------- + + +def render(plot: PlottingFn, path: pathlib.Path) -> Image.Image: + """Render a plot to a real PNG at ``path`` and load it back as a PIL image -- the + actual pixels the Critic and Judge inspect. A drawing error propagates (the + Visualizer's render-doctest already guards against non-rendering code). + + ``savefig(bbox_inches="tight")`` trims margins during the save itself -- safe on a + canvas-less OO figure, unlike a separate ``tight_layout()`` call. + """ + fig = plot() + fig.savefig(path, dpi=120, bbox_inches="tight") + return Image.open(path) + + +# --------------------------------------------------------------------------- +# Agent 1 -- the Retriever. Holds the retrieval Tool (scoped to this class), ranks +# the candidate metadata, and emits a selection certified against R. +# --------------------------------------------------------------------------- + + +class Retriever(Agent): + """You are the Retriever Agent that opens the pipeline. You perform generative + retrieval: rank the reference exemplars by how well their *visual structure* and + research domain match the task -- prioritizing diagram structure over topic + similarity -- and select the few that will best guide the downstream agents.""" + + @Tool.define + def reference_catalog(self) -> list[PlotExemplar]: + """Return the full reference set R -- every candidate exemplar's key, chart + type, domain, caption, and structure/style notes -- to rank over before + selecting.""" + return list(REFERENCES.values()) + + @Template.define + def retrieve(self, task: IllustrationTask) -> Retrieval: + """Inspect the reference set via ``reference_catalog``, then select the two or + three exemplars whose visual structure and domain best fit the task. Weight + structural/diagram-type match above topic similarity. Return their keys and a + one-line rationale; select only keys that exist in R. + + {task} + """ + + +# --------------------------------------------------------------------------- +# Agent 2 -- the Planner. Toolless: it in-context-learns from the retrieved +# exemplars and turns S + C into a structured, data-bearing description P. +# --------------------------------------------------------------------------- + + +class Planner(Agent): + """You are the Planner Agent, the cognitive core. By in-context learning from the + retrieved exemplars, you translate the source context and caption into a detailed, + structured description of the target plot -- transcribing the data exactly so the + figure will be faithful.""" + + @Template.define + def plan( + self, task: IllustrationTask, exemplars: list[PlotExemplar] + ) -> PlotDescription: + """Produce the ``PlotDescription`` for this task, learning the appropriate + chart type and composition from the retrieved exemplars. Transcribe the data + table in the source context into ``categories`` and ``series`` exactly -- every + number must come from S, in order. Fill each field as its schema describes. + + {task} + + {exemplars} + """ + + +# --------------------------------------------------------------------------- +# Agent 3 -- the Stylist. Synthesizes the aesthetic guideline G from R, then +# restyles the description P into the stylistically optimized plan P*. +# --------------------------------------------------------------------------- + + +class Stylist(Agent): + """You are the Stylist Agent, a design consultant. You first distill a reusable + aesthetic guideline from the reference set, then apply it to restyle the planner's + description into a publication-quality, stylistically optimized plan.""" + + @Template.define + def synthesize_guideline(self, exemplars: list[PlotExemplar]) -> AestheticGuideline: + """Traverse the reference exemplars' style notes and synthesize one reusable + ``AestheticGuideline`` for academic statistical plots. + + {exemplars} + """ + + @Template.define + def restyle( + self, description: PlotDescription, guideline: AestheticGuideline + ) -> StyledPlot: + """Restyle the planner's description into the optimized plan P* by bundling it + with the aesthetic guideline and writing concrete ``directives`` that apply the + guideline to this specific plot. Carry the description's data through + unchanged -- restyling never alters the numbers. + + {description} + + {guideline} + """ + + +# --------------------------------------------------------------------------- +# Agent 4 -- the Visualizer. Writes Matplotlib code (a Plot callable); the doctest +# makes "the plot must actually render" a decode-time contract. +# --------------------------------------------------------------------------- + + +class Visualizer(Agent): + """You are the Visualizer Agent, an expert Matplotlib programmer. You answer by + writing code: you turn a plan into a function that draws the plot, and the harness + renders it. You never reason the figure out in prose -- you draw it.""" + + @Template.define + def visualize(self, plan: StyledPlot) -> PlottingFn: + """Write ``plot``: a nullary function that BUILDS and RETURNS a fresh + matplotlib ``Figure`` via the object-oriented API. Inside, do: + ``fig = Figure(figsize=(8, 5)); ax = fig.subplots()``, draw onto ``ax``, and + ``return fig``. ``Figure`` is available in scope (or ``from matplotlib.figure + import Figure`` inside the function). Do NOT use ``pyplot``/``plt`` and do NOT + call ``savefig`` -- the harness saves the returned figure. + + + {plan} + + + Read ALL data and labels from the ``plan`` object, which is in scope + (``plan.description.categories``, ``plan.description.series``, etc.) -- do not + hardcode values, so the same code draws any plan. Apply the plan's palette and + style directives. + + Example usage: + + >>> _SMOKE_PLAN = StyledPlot( + ... description=PlotDescription( + ... chart_type="grouped_bar", + ... title="smoke", + ... x_label="group", + ... y_label="value", + ... categories=["p", "q"], + ... series=[DataSeries("m1", [1.0, 2.0]), DataSeries("m2", [3.0, 4.0])], + ... notes="two groups, two series", + ... ), + ... guideline=AestheticGuideline( + ... palette=["#4C72B0", "#DD8452"], + ... marks_and_containers="plain bars", + ... lines_and_arrows="none", + ... layout="grouped", + ... typography="default", + ... icons="none", + ... ), + ... directives="grouped bars, legend, y from 0", + ... ) + >>> isinstance(Visualizer().visualize(_SMOKE_PLAN)(), Figure) + True + """ + + +# --------------------------------------------------------------------------- +# Agent 5 -- the Critic. Sees the rendered PNG and refines the plan. A stateless +# Agent method (a fresh instance per loop iteration), never a module-level Template. +# --------------------------------------------------------------------------- + + +class Critic(Agent): + """You are the Critic Agent. You close the refinement loop: you look at the + actually-rendered plot, judge it against the source context and caption, and hand + the Visualizer a refined plan that fixes what you saw.""" + + @Template.define + def critique( + self, image: Image.Image, task: IllustrationTask, plan: StyledPlot + ) -> Critique: + """Here is the plot rendered from the current plan. Inspect the IMAGE against + the source context S and the caption C. Following the paper, target three + things: (1) factual misalignments -- are the numbers, categories, and labels + correct and complete vs. S and C?; (2) visual glitches -- overlap, clipping, + missing legend, unreadable text, clutter; and (3) areas for improvement -- + concrete readability/aesthetic upgrades even when nothing is strictly wrong + (clearer emphasis of the proposed method, better label/legend placement, + gridline and spine styling, value labels, headroom). List what you actually + see and return a refined plan that applies it, always keeping the data true to + S. Leave issues empty only if the plot is already publication-ready. + + {image} + + {task} + + + {plan} + + """ + + +# --------------------------------------------------------------------------- +# The iterative refinement loop -- the Visualizer<->Critic cycle, made real. +# --------------------------------------------------------------------------- + + +def refine( + task: IllustrationTask, + plan: StyledPlot, + *, + max_iter: int, + outdir: pathlib.Path, +) -> tuple[Image.Image, Image.Image, StyledPlot]: + """Run the T-round Visualizer<->Critic loop. I_0 = render(P*); each round the + Critic inspects the current render and refines the plan, which the Visualizer + re-renders (final output I_T). Returns (round-0 image, final image, final plan). + + Defaults to the paper's *fixed* T rounds -- the Critic always emits a refined + description P_{t+1}, as in the paper. With ``early_stop`` it may halt once the + Critic reports no remaining issues. + + A fresh Visualizer/Critic per iteration keeps them stateless, so each render is + judged on its own (nothing anchors on an earlier round's verdict). + """ + round0 = img = render(Visualizer().visualize(plan), outdir / "round_0.png") + for t in range(max_iter): + critique = Critic().critique(img, task, plan) + if critique.issues: + plan = critique.refined + img = render(Visualizer().visualize(plan), outdir / f"round_{t + 1}.png") + else: + break + + return round0, img, plan + + +# --------------------------------------------------------------------------- +# The pipeline -- the two phases threaded together. +# --------------------------------------------------------------------------- + + +def illustrate( + task: IllustrationTask, + *, + max_iter: int, + outdir: pathlib.Path, +) -> StyledPlot: + """Linear Planning Phase (Retriever -> Planner -> Stylist) then the Iterative + Refinement Loop (Visualizer <-> Critic), judged by the referenced ablation + round-0 vs final. Returns the final plan, the judge's per-dimension scores, and + the aggregated overall winner ('A'=round-0, 'B'=final, or 'tie').""" + # Linear Planning Phase. + retrieval = Retriever().retrieve(task) + description = Planner().plan(task, retrieval.selected) + + stylist = Stylist() + guideline = stylist.synthesize_guideline(list(REFERENCES.values())) + p_star = stylist.restyle(description, guideline) + + # Iterative Refinement Loop: the real render+critique cycle. + round0, final, plan = refine(task, p_star, max_iter=max_iter, outdir=outdir) + return plan + + +# --------------------------------------------------------------------------- +# Demo task: a grouped-bar comparison whose data lives in S, so faithfulness is crisp. +# --------------------------------------------------------------------------- + +DEMO_TASK = IllustrationTask( + source_context="""\ +We evaluate three methods -- Baseline, Ours, and Ours+Aug -- on three image +classification benchmarks, reporting top-1 accuracy (%). The measured results: + + Method | CIFAR-10 | SVHN | STL-10 + -----------+----------+-------+------- + Baseline | 71.2 | 88.4 | 64.9 + Ours | 78.5 | 91.2 | 70.3 + Ours+Aug | 82.1 | 92.8 | 73.6 + +Ours improves over Baseline on every benchmark, and adding augmentation (Ours+Aug) +improves further; the ordering Baseline < Ours < Ours+Aug holds on all three.""", + intent="Overall comparison of the three methods across the three benchmarks.", +) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--rounds", + type=int, + default=3, + help="Visualizer<->Critic refinement rounds (the paper's fixed T=3)", + ) + parser.add_argument( + "--outdir", + type=str, + default=None, + metavar="DIR", + help="Directory for rendered PNGs (defaults to a fresh temp dir)", + ) + args = parser.parse_args() + + outdir = ( + pathlib.Path(args.outdir) + if args.outdir is not None + else pathlib.Path(tempfile.mkdtemp(prefix="paperbanana_")) + ) + outdir.mkdir(parents=True, exist_ok=True) + print(f"Task caption: {DEMO_TASK.intent}") + + plan = illustrate(DEMO_TASK, max_iter=args.rounds, outdir=outdir) + + print("\n[final plan]") + print(plan) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoresearch/implementation.py b/docs/source/llm_examples/autoresearch/implementation.py new file mode 100644 index 000000000..a1b4d834c --- /dev/null +++ b/docs/source/llm_examples/autoresearch/implementation.py @@ -0,0 +1,800 @@ +"""MARS: budget-aware modular ML engineering as a cost-constrained tree search. + +Implements the core of "MARS: Modular Agent with Reflective Search for Automated +AI Research" (arXiv:2602.02660). The paper's diagnosis is that LLM coding agents +for machine-learning engineering "generate monolithic scripts that ignore +execution costs and causal factors": they write one big program in a vacuum, never +budgeting the expensive model-evaluation step and never learning *why* one attempt +beat another. Its fix rests on three pillars: (1) *budget-aware planning* via a +cost-constrained MCTS that balances solution quality against compute expense; (2) +*modular construction* via a Design -> Decompose -> Implement pipeline that breaks a +repository into modules instead of one script; and (3) *comparative reflective +memory* that compares solution paths to solve credit assignment and transfer +lessons across branches. Each pillar falls out of ordinary effectful idioms: + + * The evaluator is the ground truth, and we do not fake it. Like ``formalization`` + with the Lean compiler, MARS's load-bearing part is the *expensive model + evaluation* the paper says agents ignore, so we make it real: a synthesized + pipeline is scored by a deterministic, re-runnable Python ``evaluate`` that + returns both a quality metric (tour length) and the *measured execution cost* + of running it. The MCTS reward and the budget are therefore real numbers, not + simulated ones -- the same "ground truth a claim certifies against" role that + ``investigation``'s evaluator plays for its Solution. + + * Decompose *is* the search's action space. The Designer emits a typed ``Design`` + -- an ordered list of ``ModuleSpec``s, each carrying a few candidate + implementation *strategies* -- and those strategies are exactly the branching + actions of the MCTS tree. Coordination lives in a structured value threaded + between agents (the Outline idiom of ``writing`` / the StyledPlot idiom of + ``illustration``), and the tree is a handful of ordinary calls playing from it. + + * Implement is code synthesis with a decode-time contract. Each module is a real + ``Stage`` callable the harness compiles from the Implementer's code (the + ``Callable``-return idiom of ``investigation``'s Solver). Its docstring carries a + doctest that certifies "the module returns a valid tour (a permutation)" at + decode time -- run on a *different, smaller* instance than the evaluation one, so + a stage that hardcodes the problem size fails the doctest and is fed back by + ``RetryLLMHandler`` (the render-doctest grounding of ``illustration`` / + ``countdown``). "The module must produce a valid tour" is grounding by + construction; the Test Executor / Error Analyzer of the paper are this evaluator + plus the harness's retry feedback. + + * Budget-aware MCTS is a plain-Python loop. Selection uses the paper's + cost-aware UCT -- exploitation + exploration - alpha * (cost / budget) -- so the + search favors high-reward, low-cost branches; a rollout budget caps the search, + which is the point (enumerating every design x implementation would be far too + expensive, so the agent must *plan* which to try). Reward and cost backpropagate + up the path exactly as in textbook MCTS. + + * Comparative reflective memory is an LLM comparison feeding a refine loop. Every + few rollouts the Reflector contrasts a high-reward branch with a low-reward one + and emits a transferable ``Reflection`` (the paper's credit assignment); lessons + accumulate in a ``Memory`` spliced into later Implement prompts, and adding a + lesson invalidates the synthesis cache so subsequent branches re-implement with + it -- cross-branch transfer made observable, an LLM-judge comparison (as in + ``writing``'s reviewer) driving a re-implementation loop. + +Demonstrates: +- A real, re-runnable evaluator as ground truth: synthesized module ``Callable``s + are scored by deterministic Python (tour length) at a *measured* execution cost, + so the MCTS reward and the search budget are real, not simulated +- Budget-aware MCTS in plain Python: the paper's cost-aware UCT (exploit + explore + - alpha * cost/budget) plus a rollout budget, over a tree whose actions are the + ``Design``'s per-module strategies +- A typed ``Design`` emitted by one agent that *is* the search's action space -- + Decompose-as-data threaded through Implement (the Outline idiom) +- Code synthesis with a decode-time contract: each module is a ``Callable`` whose + doctest certifies it returns a valid permutation, fed back by ``RetryLLMHandler``, + and run on a different instance so it cannot hardcode the problem size +- Comparative reflective memory: an LLM comparison of a strong vs. weak branch + emits a lesson spliced into later syntheses (invalidating the cache), so + cross-branch transfer is observable +- Decode-time certification of the ``Design``'s shape (>= 2 strategies per module, + unique names), and per-field guidance via ``field(metadata={"description": ...})`` +""" + +# Simplifications vs. the source: +# - One planted MLE-style task, run end to end, not MLE-Bench's Kaggle repositories. +# The task is budget-constrained Euclidean TSP; the "repository" is a short +# pipeline of composed ``Stage`` functions rather than a multi-file project, and +# quality is tour length -- a stand-in for a real benchmark metric. This shows the +# Design-Decompose-Implement *shape* and the cost/quality tradeoff, not ML at scale. +# - Cost is measured wall-clock execution time of the synthesized pipeline (min over +# repeats), a real but machine- and noise-dependent proxy for MARS's "expensive +# model evaluation"; the budget is a rollout cap plus this cost feeding UCT, not a +# token-accounted API budget. Because the cost signal is real (hence noisy), which +# design wins can vary run to run -- apt for a genuine cost, but not a fixed golden +# output (contrast the deterministic-corpus examples). +# - The MCTS is small (a shallow tree of a few modules x a few strategies) and +# rollouts complete by random strategy choice rather than a learned default policy; +# there is no progressive widening. It demonstrates the cost-aware search shape. +# - Reflective memory compares the current best vs. worst successful branch every few +# rollouts and splices lessons textually; there is no embedding store and no reward +# re-weighting of tree nodes from lessons (a lesson acts only by re-implementation). +# The paper's "63% of lessons come from cross-branch transfer" is reported here only +# as a simple post-hoc count on one task, not reproduced as a statistic. +# - The Implementer writes pure Python over the given cities; there is no separate +# refactor/debug sub-loop beyond the harness's synth + doctest + retry. +# - No ContextVar: unlike ``investigation``/``formalization``, nothing certifies +# against per-run mutable state -- the module doctest checks a structural invariant +# and the evaluator is handed its stages explicitly, so ground truth stays local. + +import argparse +import collections.abc +import dataclasses +import inspect +import math +import random +import time + +import pydantic + +from effectful.handlers.llm import Agent, Template + +# A field's ``metadata={"description": ...}`` is inlined by pydantic into that +# field's JSON schema, which the harness renders into the system prompt as part of a +# template's argument (and structured-output) spec. So per-field guidance reaches the +# model *through the type* -- used below only where the field name and type don't +# already say it, so no prompt has to repeat it. + + +# --------------------------------------------------------------------------- +# The task and its evaluator -- the ground truth every candidate is scored by. This +# is the load-bearing part we do not fake: a real, re-runnable Python evaluator that +# returns both a quality metric and the *measured* execution cost of running the +# synthesized pipeline (MARS's "expensive model evaluation", made real). +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class City: + """A point in the plane the tour must visit.""" + + x: float + y: float + + +@pydantic.dataclasses.dataclass(frozen=True) +class Task: + """One MLE-style engineering task: visit every city once and return, minimizing + total Euclidean distance. A stand-in for a benchmark whose metric is expensive to + evaluate and whose best solution trades quality against compute.""" + + cities: tuple[City, ...] + + +# A tour is a permutation of city indices; the pipeline's job is to reorder it to +# shorten the round trip. A Stage is one module of the pipeline: it takes the cities +# and the current tour and returns an improved tour. Uniform typing makes the modules +# compose by a plain fold and keeps synthesis robust. +type Tour = list[int] +type Stage = collections.abc.Callable[[list[City], Tour], Tour] + + +def distance(a: City, b: City) -> float: + return math.hypot(a.x - b.x, a.y - b.y) + + +def tour_length(cities: collections.abc.Sequence[City], tour: Tour) -> float: + """Total length of the closed tour that visits ``cities`` in ``tour`` order.""" + n = len(tour) + return sum(distance(cities[tour[i]], cities[tour[(i + 1) % n]]) for i in range(n)) + + +def _validate_tour(tour: Tour, n: int) -> None: + """A stage's output must be a permutation of all ``n`` city indices, or it is not + a valid tour -- the invariant the module doctest also enforces at decode time.""" + if sorted(tour) != list(range(n)): + raise ValueError( + f"a stage returned {tour}, which is not a valid tour: it must be a " + f"permutation of every city index 0..{n - 1} exactly once" + ) + + +def run_pipeline(cities: collections.abc.Sequence[City], stages: list[Stage]) -> Tour: + """Fold the identity tour through every stage, certifying each stage's output is a + valid permutation. A stage that returns garbage raises -- the same + certification-by-construction the doctest makes at decode time.""" + tour: Tour = list(range(len(cities))) + for stage in stages: + tour = list(stage(list(cities), tour)) + _validate_tour(tour, len(cities)) + return tour + + +# Repeat the pipeline a few times and take the minimum runtime: the standard robust +# estimator for a small computation's cost, damping OS/scheduler noise. +COST_REPEATS = 5 + + +def evaluate(task: Task, stages: list[Stage]) -> tuple[float, float]: + """Run the assembled pipeline and return ``(tour_length, cost_seconds)`` -- the + real quality metric and the measured execution cost. Both feed the MCTS: length + becomes the reward, cost enters cost-aware UCT and the budget. Raises (via + ``run_pipeline``) if any stage produces an invalid tour.""" + cities = list(task.cities) + best_cost = math.inf + tour: Tour = list(range(len(cities))) + for _ in range(COST_REPEATS): + start = time.perf_counter() + tour = run_pipeline(cities, stages) + best_cost = min(best_cost, time.perf_counter() - start) + return tour_length(cities, tour), best_cost + + +# --------------------------------------------------------------------------- +# Structured artifacts crossing between agents. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class ModuleSpec: + """One stage of the pipeline the Designer decomposes the task into: what it should + accomplish, plus the candidate implementation *strategies* that become the MCTS + branching actions for this stage.""" + + name: str + intent: str = dataclasses.field( + metadata={ + "description": "What this stage does to the tour it receives (e.g. build " + "an initial ordering, or locally improve the incoming tour). Every stage " + "takes the cities and the current tour and returns a valid tour." + } + ) + strategies: list[str] = dataclasses.field( + metadata={ + "description": "Two or three distinct, concrete implementation approaches " + "for this stage (e.g. 'nearest-neighbour construction', '2-opt local " + "search', 'or-opt segment moves'). Each becomes one search action." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Design: + """The Design/Decompose artifact: the pipeline's architecture as an ordered list + of modules. This one structured value *is* the MCTS action space -- every + downstream Implement call and every tree action reads it.""" + + analysis: str = dataclasses.field( + metadata={ + "description": "A short reading of the task: what makes a good tour and " + "how the modules cooperate to produce one." + } + ) + modules: list[ModuleSpec] + + def __post_init__(self) -> None: + if not self.modules: + raise ValueError("a design must have at least one module (pipeline stage)") + names = [m.name for m in self.modules] + if len(set(names)) != len(names): + raise ValueError(f"module names must be unique, got {names}") + for m in self.modules: + if len(set(m.strategies)) < 2: + raise ValueError( + f"module {m.name!r} must offer at least two distinct strategies " + f"(the search needs branching actions), got {m.strategies}" + ) + + def __str__(self) -> str: + lines = [f"analysis: {self.analysis}"] + for i, m in enumerate(self.modules): + lines.append(f" module {i} [{m.name}]: {m.intent}") + lines += [f" - {s}" for s in m.strategies] + return "\n".join(lines) + + +@pydantic.dataclasses.dataclass(frozen=True) +class PipelineChoice: + """One decided stage of a full pipeline: which module, and the strategy chosen for + it. A list of these is the path the MCTS committed to.""" + + module: str + strategy: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class RolloutSummary: + """A finished branch handed to the Reflector: which strategies it chose, and the + real outcome the evaluator measured. The Reflector compares two of these to assign + credit.""" + + choices: list[PipelineChoice] + tour_length: float + cost_seconds: float + reward: float = dataclasses.field( + metadata={ + "description": "The search reward: fractional improvement of the tour over " + "the trivial identity ordering, in [0, 1] (higher is better)." + } + ) + + def __str__(self) -> str: + picks = " -> ".join(f"{c.module}:{c.strategy}" for c in self.choices) + return ( + f"[{picks}] length={self.tour_length:.1f} " + f"cost={self.cost_seconds * 1e3:.2f}ms reward={self.reward:.3f}" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Reflection: + """The Reflector's credit-assignment output: one transferable lesson drawn from + comparing a strong branch against a weak one.""" + + lesson: str = dataclasses.field( + metadata={ + "description": "A concrete, transferable engineering lesson about how to " + "implement a stage better -- grounded in the difference between the two " + "branches, not generic advice. It will be shown to future implementers." + } + ) + applies_to: str = dataclasses.field( + metadata={ + "description": "Which module or strategy this lesson informs, so a future " + "implementer knows when it is relevant." + } + ) + + def __str__(self) -> str: + return f"({self.applies_to}) {self.lesson}" + + +@dataclasses.dataclass +class Memory: + """The comparative reflective memory: the accumulated lessons, spliced into later + Implement prompts. Not frozen -- reflections are appended as the search learns.""" + + reflections: list[Reflection] = dataclasses.field(default_factory=list) + + def digest(self) -> str: + """The lessons rendered for an Implement prompt; empty guidance when none.""" + if not self.reflections: + return "(no lessons learned yet)" + return "\n".join(f"- {r}" for r in self.reflections) + + +# --------------------------------------------------------------------------- +# The three agents. All are closed-book: no search tools -- the only "tool" is code +# synthesis (the harness's FinalTool) and the deterministic evaluator. This is the +# distinctive shape of a coding agent, versus the literature examples' search tools. +# --------------------------------------------------------------------------- + + +class Designer(Agent): + """You are the Design & Decompose agent that opens the pipeline. Instead of writing + one monolithic script, you break the task into a short, ordered pipeline of modules + (stages), and for each module you propose a few concrete implementation strategies + for a downstream search to choose among.""" + + @Template.define + def design(self, task: Task) -> Design: + """Analyze the task and decompose a solution into an ordered pipeline of two + or three modules. Each module is a stage that takes the cities and the current + tour and returns an improved, valid tour; a natural decomposition is + construction (build an initial tour) followed by one or more local-improvement + stages. For each module, propose two or three *distinct* implementation + strategies -- these become the choices a budget-aware search explores. Fill + each field as its schema describes. + + {task} + """ + + +class Implementer(Agent): + """You are the Implement agent, an expert Python programmer. You answer by writing + code, not prose: you turn one module of the design into a function that transforms + a tour, and the harness compiles and runs it. You read the module's intent and the + chosen strategy, and you apply any lessons learned from earlier attempts.""" + + @Template.define + def implement(self, module: ModuleSpec, strategy: str, lessons: str) -> Stage: + """Write ``stage``: a function ``stage(cities, tour)`` that takes the list of + ``City`` points and the current ``tour`` (a list of city indices) and RETURNS + an improved tour -- a list containing every index ``0..len(cities)-1`` exactly + once. Implement the module's intent using the chosen strategy. + + Read everything from the ``cities`` and ``tour`` arguments (use + ``len(cities)`` for the size, ``city.x`` / ``city.y`` for coordinates, and + ``math`` if you need it) -- do NOT hardcode the number of cities or any + coordinates, so the same code works for any instance. Return a valid + permutation; never drop, duplicate, or invent an index. + + Module: {module.name} -- {module.intent} + Strategy to implement: {strategy} + + Lessons learned from earlier attempts (apply any that are relevant): + {lessons} + + The doctest runs the synthesized stage on a tiny four-city instance, while + the real evaluation runs it on the demo task (much larger), so a stage that + hardcodes the problem size fails the doctest and is corrected -- the + anti-hardcode trick of ``illustration``. The recursive + ``Implementer().implement`` call is routed to your own submission. + + >>> _module = ModuleSpec( + ... name="reorder", + ... intent="reorder the incoming tour to shorten the round trip", + ... strategies=["greedy nearest-neighbour", "swap crossing edges"], + ... ) + >>> _cities = [City(0.0, 0.0), City(1.0, 0.0), City(1.0, 1.0), City(0.0, 1.0)] + >>> _stage = Implementer().implement(_module, "greedy nearest-neighbour", "") + >>> _out = _stage(_cities, [0, 1, 2, 3]) + >>> sorted(_out) == [0, 1, 2, 3] + True + """ + + +class Reflector(Agent): + """You are the Comparative Reflection agent. You look at two finished branches -- + one that scored well and one that scored poorly -- and you diagnose *why* the good + one won, distilling a single transferable lesson a future implementer can reuse. + You solve credit assignment by comparison, not by guessing.""" + + @Template.define + def reflect(self, better: RolloutSummary, worse: RolloutSummary) -> Reflection: + """Compare these two branches of the search. The first achieved a higher reward + (a shorter tour, accounting for its execution cost) than the second. Identify + the concrete difference in their strategy choices or implementation that most + plausibly explains the gap, and state one transferable lesson for implementing + such a stage better next time. Ground the lesson in the comparison -- what the + better branch did that the worse one did not -- not in generic advice. + + {better} + + {worse} + """ + + +# --------------------------------------------------------------------------- +# Budget-aware MCTS. The tree's actions are the Design's per-module strategies; a +# leaf is a full pipeline, scored by the real evaluator. Cost-aware UCT and a rollout +# budget make the search prefer high-quality, low-cost designs. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Node: + """One MCTS node: a partial pipeline. ``path`` is the strategy chosen for each + module decided so far (module ``len(path)`` is decided by this node's children). + ``untried`` holds the strategies for the next module not yet expanded.""" + + path: tuple[str, ...] + untried: list[str] + children: list["Node"] = dataclasses.field(default_factory=list) + visits: int = 0 + reward_sum: float = 0.0 + cost_sum: float = 0.0 + + @property + def avg_reward(self) -> float: + return self.reward_sum / self.visits if self.visits else 0.0 + + @property + def avg_cost(self) -> float: + return self.cost_sum / self.visits if self.visits else 0.0 + + +@dataclasses.dataclass +class Rollout: + """The record of one evaluated pipeline: its choices, the compiled stages, and the + real outcome. ``ok`` is False if synthesis or evaluation failed (that branch scores + nothing). ``lessons_available`` records how many reflections existed when it ran -- + used to measure cross-branch transfer afterwards.""" + + path: tuple[str, ...] + stages: list[Stage] + length: float + cost: float + reward: float + ok: bool + lessons_available: int + + +@dataclasses.dataclass +class Search: + """A budget-aware MCTS over the Design's action space, with comparative reflective + memory. Holds everything one task's search threads together, so no ambient/global + state is needed (contrast the ContextVar examples): the evaluator is handed its + stages explicitly and the doctest certifies a structural invariant.""" + + task: Task + design: Design + exploration: float # UCT exploration constant C + cost_weight: float # UCT cost coefficient alpha + rng: random.Random + memory: Memory = dataclasses.field(default_factory=Memory) + root: Node = dataclasses.field(init=False) + # A synthesized stage is expensive to produce, so cache it by (generation, module + # index, strategy). Adding a lesson bumps ``generation``, invalidating the cache so + # later branches re-implement with the lesson -- how cross-branch transfer bites. + _cache: dict[tuple[int, int, str], Stage] = dataclasses.field(default_factory=dict) + generation: int = 0 + baseline: float = 0.0 # length of the identity tour -- the reward's zero point + max_cost: float = 1e-9 # running max rollout cost, to normalize the UCT penalty + rollouts: list[Rollout] = dataclasses.field(default_factory=list) + + def __post_init__(self) -> None: + self.root = Node(path=(), untried=list(self.design.modules[0].strategies)) + self.baseline = tour_length( + self.task.cities, list(range(len(self.task.cities))) + ) + + # --- reward ----------------------------------------------------------- + + def reward_of(self, length: float) -> float: + """Fractional improvement of a tour over the identity ordering, clamped to + [0, 1] -- a bounded reward so the UCT exploration term stays well-scaled.""" + return max(0.0, min(1.0, (self.baseline - length) / self.baseline)) + + def score(self, reward: float, cost: float) -> float: + """The budget-aware value of a branch: reward minus a normalized cost penalty. + This is the UCT *exploitation* term and the final selection key -- the paper's + 'favor high-reward, low-cost branches'.""" + return reward - self.cost_weight * (cost / self.max_cost) + + def uct(self, child: Node, parent: Node) -> float: + """Cost-aware UCT: exploitation + exploration - alpha * cost/budget (the + paper's selection criterion). Unvisited children sort first.""" + if child.visits == 0: + return math.inf + explore = self.exploration * math.sqrt(math.log(parent.visits) / child.visits) + return self.score(child.avg_reward, child.avg_cost) + explore + + # --- tree policy ------------------------------------------------------ + + def select(self) -> list[Node]: + """Descend from the root by cost-aware UCT until reaching a node that can be + expanded (has untried strategies) or is terminal (a full pipeline). Returns the + path of nodes visited, for backpropagation.""" + node = self.root + path = [node] + while not node.untried and node.children: # fully expanded, non-terminal + node = max(node.children, key=lambda c: self.uct(c, node)) + path.append(node) + return path + + def expand(self, node: Node) -> Node: + """Add one child for an untried strategy of the next module.""" + strategy = node.untried.pop(0) + depth = len(node.path) + 1 + untried = ( + list(self.design.modules[depth].strategies) + if depth < len(self.design.modules) + else [] + ) + child = Node(path=node.path + (strategy,), untried=untried) + node.children.append(child) + return child + + def complete(self, path: tuple[str, ...]) -> tuple[str, ...]: + """Finish a partial path into a full pipeline by choosing a random strategy for + each remaining module (the default rollout policy).""" + full = list(path) + full.extend( + self.rng.choice(self.design.modules[depth].strategies) + for depth in range(len(path), len(self.design.modules)) + ) + return tuple(full) + + # --- rollout ---------------------------------------------------------- + + def build(self, path: tuple[str, ...]) -> list[Stage]: + """Synthesize (or reuse) the stage for each chosen module. Caching keys on the + current ``generation`` so a new lesson forces re-implementation.""" + stages: list[Stage] = [] + for depth, strategy in enumerate(path): + key = (self.generation, depth, strategy) + if key not in self._cache: + self._cache[key] = Implementer().implement( + self.design.modules[depth], strategy, self.memory.digest() + ) + stages.append(self._cache[key]) + return stages + + def rollout(self, node: Node) -> Rollout: + """Complete the node's path to a full pipeline, synthesize it, and evaluate -- + the real quality and cost. A synthesis or evaluation failure scores nothing.""" + full = self.complete(node.path) + try: + stages = self.build(full) + length, cost = evaluate(self.task, stages) + reward, ok = self.reward_of(length), True + except Exception as exc: # retries exhausted, or an invalid-tour stage + print(f" [rollout] {full} failed: {type(exc).__name__}") + stages, length, cost, reward, ok = [], math.inf, 0.0, 0.0, False + self.max_cost = max(self.max_cost, cost) + r = Rollout( + path=full, + stages=stages, + length=length, + cost=cost, + reward=reward, + ok=ok, + lessons_available=len(self.memory.reflections), + ) + self.rollouts.append(r) + return r + + def backpropagate(self, path: list[Node], reward: float, cost: float) -> None: + for node in path: + node.visits += 1 + node.reward_sum += reward + node.cost_sum += cost + + # --- comparative reflection ------------------------------------------ + + def reflect(self) -> None: + """Compare the best and worst distinct successful branches so far and store a + lesson, then bump the generation so later branches re-implement with it. This + is the paper's cross-path credit assignment feeding the reflective memory.""" + ok = [r for r in self.rollouts if r.ok] + if len(ok) < 2: + return + best = max(ok, key=lambda r: r.reward) + worst = min(ok, key=lambda r: r.reward) + if best.path == worst.path: + return + reflection = Reflector().reflect(self._summ(best), self._summ(worst)) + self.memory.reflections.append(reflection) + self.generation += ( + 1 # invalidate the synthesis cache: re-implement with the lesson + ) + print(f" [reflect] lesson: {reflection}") + + def _summ(self, r: Rollout) -> RolloutSummary: + choices = [ + PipelineChoice(self.design.modules[d].name, s) for d, s in enumerate(r.path) + ] + return RolloutSummary(choices, r.length, r.cost, r.reward) + + # --- driver ----------------------------------------------------------- + + def run(self, *, max_rollouts: int, reflect_every: int) -> Rollout: + """The MCTS loop under a rollout budget: select -> expand -> rollout -> + backpropagate, reflecting every few rollouts. Returns the best actually- + evaluated pipeline (balancing quality and cost) -- MARS's best-path extraction.""" + for i in range(1, max_rollouts + 1): + path = self.select() + leaf = path[-1] + if leaf.untried: # expand a new action + leaf = self.expand(leaf) + path.append(leaf) + result = self.rollout(leaf) + self.backpropagate(path, result.reward, result.cost) + print( + f" rollout {i}/{max_rollouts}: {self._summ(result)}" + if result.ok + else f" rollout {i}/{max_rollouts}: (failed)" + ) + if reflect_every and i % reflect_every == 0: + self.reflect() + + succeeded = [r for r in self.rollouts if r.ok] + if not succeeded: + raise RuntimeError("every rollout failed to produce a valid pipeline") + return max(succeeded, key=lambda r: self.score(r.reward, r.cost)) + + def cross_branch_gain(self) -> tuple[int, float, float]: + """A simple post-hoc read on whether lessons helped later branches: the best + reward reached *before* any lesson existed, versus how many later branches beat + it. A nod to the paper's cross-branch-transfer analysis, not its statistic.""" + pre = [r.reward for r in self.rollouts if r.ok and r.lessons_available == 0] + post = [r.reward for r in self.rollouts if r.ok and r.lessons_available > 0] + best_pre = max(pre, default=0.0) + best_post = max(post, default=0.0) + improved = sum(1 for r in post if r > best_pre) + return improved, best_pre, best_post + + +# --------------------------------------------------------------------------- +# The pipeline: Design -> (budget-aware MCTS over Decompose/Implement) with reflection. +# --------------------------------------------------------------------------- + + +def implement( + task: Task, + *, + max_rollouts: int, + reflect_every: int, + exploration: float, + cost_weight: float, + seed: int, +) -> tuple[Search, Rollout]: + """Design the pipeline, then run the cost-constrained MCTS over its modules and + strategies -- synthesizing and evaluating each explored pipeline, and reflecting + across branches -- and return the search and the best pipeline found.""" + print("[design] decomposing the task into a modular pipeline ...") + design = Designer().design(task) + print(design) + + search = Search( + task=task, + design=design, + exploration=exploration, + cost_weight=cost_weight, + rng=random.Random(seed), + ) + print( + f"\n[search] budget-aware MCTS: {max_rollouts} rollouts, " + f"baseline tour length {search.baseline:.1f}\n" + ) + best = search.run(max_rollouts=max_rollouts, reflect_every=reflect_every) + return search, best + + +# --------------------------------------------------------------------------- +# Demo task: a planted set of cities. Generated deterministically from a seed so the +# instance is fixed, while the search (and its real, noisy cost signal) does the work. +# --------------------------------------------------------------------------- + + +def make_task(num_cities: int, seed: int) -> Task: + rng = random.Random(seed) + cities = tuple( + City(rng.uniform(0, 100), rng.uniform(0, 100)) for _ in range(num_cities) + ) + return Task(cities=cities) + + +def _print_stage_source(stages: list[Stage]) -> None: + """Show the code MARS actually wrote for the winning pipeline, when the synthesized + source is recoverable (the eval provider registers it with ``linecache``).""" + for i, stage in enumerate(stages): + try: + src = inspect.getsource(stage) + except (OSError, TypeError): + continue + print(f"\n--- stage {i} ---\n{src.rstrip()}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--num-cities", type=int, default=40, help="Cities in the planted TSP task" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Seed for the task and rollout policy" + ) + parser.add_argument( + "--max-rollouts", + type=int, + default=6, + help="Rollout budget for the cost-aware MCTS", + ) + parser.add_argument( + "--reflect-every", + type=int, + default=3, + help="Run comparative reflection every N rollouts (0 disables it)", + ) + parser.add_argument( + "--exploration", + type=float, + default=1.4, + help="UCT exploration constant C", + ) + parser.add_argument( + "--cost-weight", + type=float, + default=0.3, + help="UCT cost coefficient alpha (how much execution cost is penalized)", + ) + args = parser.parse_args() + + task = make_task(args.num_cities, args.seed) + print(f"Task: shortest closed tour over {args.num_cities} cities\n") + + search, best = implement( + task, + max_rollouts=args.max_rollouts, + reflect_every=args.reflect_every, + exploration=args.exploration, + cost_weight=args.cost_weight, + seed=args.seed, + ) + + print("\n" + "=" * 72) + summary = search._summ(best) + print(f"Best pipeline: {summary}") + print( + f" improvement over baseline: " + f"{(search.baseline - best.length) / search.baseline * 100:.1f}%" + ) + + if search.memory.reflections: + print("\nLessons learned (comparative reflective memory):") + for r in search.memory.reflections: + print(f" - {r}") + improved, best_pre, best_post = search.cross_branch_gain() + print( + f"\nCross-branch transfer: best reward before any lesson {best_pre:.3f}; " + f"{improved} later branch(es) beat it (best after {best_post:.3f})." + ) + + _print_stage_source(best.stages) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoresearch/investigation.py b/docs/source/llm_examples/autoresearch/investigation.py new file mode 100644 index 000000000..db0e70056 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/investigation.py @@ -0,0 +1,885 @@ +"""ScientistOne: verifiable autonomous research via Chain-of-Evidence. + +Implements the core of "ScientistOne: Towards Human-Level Autonomous Research via +Chain-of-Evidence" (arXiv:2605.26340). The paper's observation is that autonomous +research agents produce professional-looking manuscripts riddled with +verifiability failures -- fabricated citations, unreproducible scores, and method +descriptions that diverge from the code -- and its fix is to make every claim +*traceable to its evidence source* rather than caught after the fact. Two +mechanisms carry that idea, and both fall out of ordinary effectful idioms: + + * Chain-of-Evidence *by construction*. Every value the Writer emits is a Claim + that certifies itself against a ground-truth Workspace at decode time. A + hallucinated citation or an invented score raises during decoding, and the + harness's ``RetryLLMHandler`` feeds the error back so the Writer must ground + the claim before it stands -- exactly the retry path ``error_recovery.py`` + uses for a bad ``Rating``. This is why ScientistOne reports zero hallucinated + references: an ungrounded reference is simply not a well-typed Claim. + + * The post-hoc CoE Audit. Four integrity checks -- I1 score verification, I2 + specification violation, I3 reference verification, I4 method-code alignment + -- run over the finished paper *uniformly*, the same way you would audit a + baseline that has no provenance of its own. Some are deterministic Python + (re-run the evaluator; re-resolve every bibkey) and some are majority-vote LLM + judges (does the code cheat? does the method match it? does each reference + actually support its claim?), like the reviewer in ``research_agent.py``. + +Demonstrates: +- Decode-time certification of structured output against external ground truth, + so ``RetryLLMHandler`` turns fabrications into corrections (Chain-of-Evidence) +- Multi-hop evidence chains: a ``ConclusionClaim`` rests on other claims (their + bibkeys/metrics), which rest on artifacts -- the *chain* in Chain-of-Evidence +- The three-stage pipeline: literature grounding -> discovery -> paper writing, + where writing is a critique/revise coherence loop (as in ``research_agent.py``) + layered on top of decode-time grounding -- the paper's Ground + Critic/Resolve +- Grounded literature review: the Investigator retrieves over a reference corpus + via a tool (as in ``rag.py``), filters out distractors across a draft/revise + pass on one stateful ``Agent``, and emits a ``Brief`` whose cited keys certify + against the database -- Chain-of-Evidence extended to the literature-review stage +- Parallel Explore-Exploit discovery: an ``asyncio`` fan-out (as in + ``map_reduce.py``) runs several solver branches per round, each a ``Template`` + returning a ``Callable`` the plain-Python evaluator scores, keeping the best -- + so the reported score has a real, re-runnable experiment log behind it +- A post-hoc audit mixing deterministic checks with majority-vote ``Template`` LLM + judges (each judge run several times in parallel, as in ``map_reduce.py``, and + the majority taken), applied uniformly to the finished artifact bundle +""" + +import argparse +import asyncio +import collections.abc +import contextvars +import dataclasses + +import pydantic.dataclasses + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# The research task and its canonical evaluator (the ground truth) +# --------------------------------------------------------------------------- + +SPEC = ( + "TASK: given a list of positive integers `numbers` and an integer `target`, " + "return a subset of `numbers` (each element used at most as many times as it " + "appears) whose sum is as large as possible without exceeding `target`. " + "SCORE: the achieved subset sum; higher is better. A subset that reuses an " + "unavailable number or exceeds the target scores nothing (it is invalid)." +) + +type Solution = collections.abc.Callable[ + [collections.abc.Sequence[int], int], list[int] +] + + +@pydantic.dataclasses.dataclass(frozen=True) +class Task: + numbers: tuple[int, ...] + target: int + + +type Evaluator = collections.abc.Callable[[Task, Solution], float] + + +def evaluate(task: Task, solve: Solution) -> float: + """Canonical evaluator: run a solution and return its score. + + Deterministic and re-runnable -- this is the ground truth that Stage 2 records + and that the audit's Score Verification independently re-derives. A malformed + subset raises, so a broken solver is fed its own error and revises (the same + retry path a fabricated claim takes). + """ + subset = list(solve(task.numbers, task.target)) + pool = list(task.numbers) + for n in subset: + if n not in pool: + raise ValueError(f"solution used {n}, which is not available in {pool}") + pool.remove(n) # each occurrence may be spent only once + total = sum(subset) + if total > task.target: + raise ValueError( + f"subset {subset} sums to {total}, exceeding target {task.target}" + ) + return float(total) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Reference: + """A bibliography entry: a citation key, the full citation text, and the abstract.""" + + key: str + citation: str + abstract: str + + +# The "literature" database: real references the Investigator searches over. Some +# bear directly on the task (subset-sum / knapsack / dynamic programming); the rest +# are plausible distractors, so selecting the relevant ones is genuine filtering and +# citing a real key is a real constraint rather than a foregone conclusion. +REFERENCES: list[Reference] = [ + Reference( + "bellman1957", + "Bellman, R. (1957). Dynamic Programming. Princeton University Press.", + "Introduces dynamic programming: solving a multistage optimization by " + "combining solutions to overlapping subproblems via the principle of optimality.", + ), + Reference( + "karp1972", + "Karp, R. (1972). Reducibility Among Combinatorial Problems.", + "Proves NP-completeness of 21 combinatorial problems, including knapsack and " + "subset sum, by polynomial-time reductions.", + ), + Reference( + "martello1990", + "Martello, S. & Toth, P. (1990). Knapsack Problems: Algorithms and Computer Implementations. Wiley.", + "A comprehensive treatment of exact and approximate algorithms for 0/1 knapsack, " + "subset sum, and bounded/unbounded variants.", + ), + Reference( + "pisinger1999", + "Pisinger, D. (1999). Linear Time Algorithms for Knapsack Problems with Bounded Weights. J. Algorithms.", + "Gives efficient dynamic-programming algorithms for knapsack and subset-sum " + "instances whose item weights are bounded.", + ), + Reference( + "horowitz1974", + "Horowitz, E. & Sahni, S. (1974). Computing Partitions with Applications to the Knapsack Problem. JACM.", + "The meet-in-the-middle technique: enumerate subset sums of each half and combine " + "them, solving subset sum in O(2^(n/2)) time.", + ), + Reference( + "ibarra1975", + "Ibarra, O. & Kim, C. (1975). Fast Approximation Algorithms for the Knapsack and Sum of Subset Problems. JACM.", + "A fully polynomial-time approximation scheme for knapsack and subset sum via " + "scaling and rounding of item values.", + ), + Reference( + "garey1979", + "Garey, M. & Johnson, D. (1979). Computers and Intractability. Freeman.", + "The standard reference on NP-completeness, including weak NP-hardness and " + "pseudo-polynomial dynamic programming for number problems like subset sum.", + ), + # --- distractors: real, well-known, but not about subset sum / knapsack --- + Reference( + "dijkstra1959", + "Dijkstra, E. (1959). A Note on Two Problems in Connexion with Graphs. Numerische Mathematik.", + "An efficient algorithm for single-source shortest paths in a graph with " + "non-negative edge weights.", + ), + Reference( + "rivest1978", + "Rivest, R., Shamir, A. & Adleman, L. (1978). A Method for Obtaining Digital Signatures. CACM.", + "The RSA public-key cryptosystem, based on the difficulty of factoring large " + "integers.", + ), + Reference( + "cook1971", + "Cook, S. (1971). The Complexity of Theorem-Proving Procedures. STOC.", + "Introduces NP-completeness and proves that boolean satisfiability (SAT) is " + "NP-complete.", + ), + Reference( + "vaswani2017", + "Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS.", + "The Transformer architecture, replacing recurrence with self-attention for " + "sequence transduction.", + ), + Reference( + "shannon1948", + "Shannon, C. (1948). A Mathematical Theory of Communication. Bell System Technical Journal.", + "Founds information theory: entropy, channel capacity, and the limits of " + "reliable communication.", + ), + Reference( + "knuth1998", + "Knuth, D. (1998). The Art of Computer Programming, Vol. 3: Sorting and Searching. Addison-Wesley.", + "A definitive treatment of comparison sorting, searching, and related data " + "structures.", + ), + Reference( + "lamport1978", + "Lamport, L. (1978). Time, Clocks, and the Ordering of Events in a Distributed System. CACM.", + "Logical clocks and the happens-before relation for ordering events in a " + "distributed system.", + ), +] + + +# --------------------------------------------------------------------------- +# Workspace: the artifact bundle every claim must trace back to +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Workspace: + """The evidence bundle every claim must trace back to. Holds only serializable + artifacts -- the reference database and an append-only log of recorded scores + (for score verification) -- so a claim can certify against it and a tool may + safely surface any of it to the model. The discovered solution *callable* is + held by the ``Writer`` and passed to the audit, not stored here.""" + + references: list[Reference] + log: dict[str, float] = dataclasses.field(default_factory=dict) + + +# The evidence bundle for the run currently in scope. Claim.__post_init__ has no +# parameters, so certification reaches ground truth ambiently -- but through a +# ContextVar rather than a bare global, so the binding is scoped to the pipeline +# (set/reset in ``run_scientist_one``) and safe under the concurrent template +# calls these examples make. +WORKSPACE: contextvars.ContextVar[Workspace] = contextvars.ContextVar("WORKSPACE") + + +# --------------------------------------------------------------------------- +# Chain-of-Evidence: claims that certify themselves against the Workspace +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class CitationClaim: + """A background statement supported by a cited reference. `bibkey` must be the + key of a real reference in the database.""" + + statement: str + bibkey: str + + def __post_init__(self) -> None: + known = {r.key for r in WORKSPACE.get().references} + if self.bibkey not in known: + raise ValueError( + f"citation {self.bibkey!r} does not resolve to any known reference " + f"(available keys: {sorted(known)}); cite only real works" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class NumericalClaim: + """A reported quantitative result: `value` must match the value recorded for + `metric` in the experiment log.""" + + metric: str + value: float + + def __post_init__(self) -> None: + log = WORKSPACE.get().log + recorded = log.get(self.metric) + if recorded is None: + raise ValueError( + f"metric {self.metric!r} was never measured " + f"(recorded metrics: {sorted(log)}); report only measured values" + ) + if abs(recorded - self.value) > 1e-9: + raise ValueError( + f"reported {self.metric}={self.value} but the experiment log records " + f"{recorded}; report the value the evaluator actually produced" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class MethodClaim: + """A prose description of how the discovered solution works.""" + + description: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class ConclusionClaim: + """A takeaway that builds on other claims rather than directly on an artifact.""" + + statement: str + supported_by: list[str] = dataclasses.field( + metadata={ + "description": "bibkeys and/or metrics this conclusion builds on; " + "each must already be cited or measured" + } + ) + + def __post_init__(self) -> None: + ws = WORKSPACE.get() + grounded = {r.key for r in ws.references} | set(ws.log) + if not self.supported_by: + raise ValueError("a conclusion must rest on at least one supporting claim") + dangling = [s for s in self.supported_by if s not in grounded] + if dangling: + raise ValueError( + f"conclusion rests on unverifiable supports {dangling}; every entry in " + f"supported_by must be a cited reference key or a recorded metric " + f"(available: {sorted(grounded)}); a conclusion may not introduce new evidence" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Paper: + """A research paper as structured, evidence-bound claims.""" + + title: str + background: list[CitationClaim] + results: list[NumericalClaim] + method: MethodClaim + conclusions: list[ConclusionClaim] + + def __str__(self) -> str: + """Render the structured, evidence-bound claims to prose -- provenance first, + prose last.""" + lines = [f"# {self.title}", "", "## Background"] + lines += [f"- {c.statement} [{c.bibkey}]" for c in self.background] + lines += ["", "## Method", self.method.description, "", "## Results"] + lines += [f"- {c.metric} = {c.value}" for c in self.results] + lines += ["", "## Conclusions"] + lines += [ + f"- {c.statement} (from {', '.join(c.supported_by)})" + for c in self.conclusions + ] + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Stage 1: literature grounding -- the Problem Investigator +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class Brief: + """A research brief: a problem framing plus the reference keys it relies on.""" + + summary: str + cited: list[str] = dataclasses.field( + metadata={"description": "citation keys of the references this brief relies on"} + ) + + def __post_init__(self) -> None: + known = {r.key for r in WORKSPACE.get().references} + dangling = [k for k in self.cited if k not in known] + if dangling: + raise ValueError( + f"brief cites unknown references {dangling}; cite only real works " + f"from the database (available keys: {sorted(known)})" + ) + + +def _relevance(query: str, ref: Reference) -> int: + """Keyword overlap between a query and a reference's citation + abstract.""" + text = f"{ref.citation} {ref.abstract}".lower() + return sum(text.count(term) for term in query.lower().split()) + + +class Investigator(Agent): + """You are the Problem Investigator opening an autonomous research project. You + survey the available literature and frame the problem -- what kind of task it is + and which known results bear on it -- before any solution is attempted.""" + + @Tool.define + def search_literature(self, query: str, top_k: int = 5) -> list[Reference]: + """Search the literature database by keyword and return the most relevant + references (citation and abstract). Issue several queries with different + keywords to survey the field before committing to a brief.""" + refs = WORKSPACE.get().references + scored = [(_relevance(query, r), r) for r in refs] + hits = sorted( + (sr for sr in scored if sr[0] > 0), reverse=True, key=lambda sr: sr[0] + ) + return [r for _, r in hits[:top_k]] + + @Template.define + def investigate(self, spec: str) -> Brief: + """Survey the literature and produce a research brief for the task below. + Use search_literature to find relevant prior work -- issue a few queries + with different keywords, read the abstracts, and decide which references + genuinely bear on this task, ignoring unrelated ones. Frame the problem in + a few sentences that refer to the relevant works by citation key. + + Task specification: + {spec} + """ + + +# --------------------------------------------------------------------------- +# Stage 2: discovery -- a parallel explore-exploit search over solver branches +# --------------------------------------------------------------------------- + +# Distinct angles for the parallel branches to pursue (explore). +APPROACHES: list[str] = [ + "an exact dynamic program over reachable subset sums", + "a greedy construction refined by local search / swaps", + "meet-in-the-middle: enumerate half-subset sums and combine", +] + + +class Solver(Agent): + """You are a careful algorithm designer and expert Python programmer. You + answer by writing code, not prose: you implement the solution as a function + and let the evaluator judge it.""" + + @Template.define + def discover( + self, spec: str, brief: str, approach: str, incumbent: float + ) -> Solution: + """Implement a solution to the task by writing ``solve``; annotate its + parameters and return type (the harness needs the annotations to compile + it). Do not read or hardcode against any particular test input. + + Pursue this approach: {approach} + Best valid score any branch has reached so far: {incumbent} -- aim to beat it. + + Task specification: + {spec} + + Research brief: + {brief} + """ + + +async def discover_best( + task: Task, brief: str, *, rounds: int, branches: int +) -> tuple[Solution, float]: + """Parallel Explore-Exploit discovery: each round runs several isolated solver + branches concurrently (explore, one approach each), scores every candidate on + the canonical evaluator, and keeps the best across rounds (best-run selection). + An invalid solution -- one that ``evaluate`` rejects -- scores nothing, which is + how spec-violating candidates are filtered out. The incumbent score is fed to + the next round so branches try to beat it (exploit). + """ + approaches = APPROACHES[:branches] + best: tuple[Solution, float] | None = None + + for r in range(rounds): + # The empty subset always scores 0, so 0.0 is the floor to beat in round 0. + incumbent = best[1] if best is not None else 0.0 + + async def branch(approach: str) -> tuple[str, Solution | None, float]: + # A fresh Solver per branch = an isolated solver cycle with its own history. + # A branch that fails to synthesize a valid, runnable solution scores + # nothing and is dropped -- best-run selection filters it out. + try: + solve = await asyncio.to_thread( + Solver().discover, SPEC, brief, approach, incumbent + ) + return approach, solve, await asyncio.to_thread(evaluate, task, solve) + except Exception: + return approach, None, 0.0 + + for approach, solve, score in await asyncio.gather( + *(branch(a) for a in approaches) + ): + if solve is None: + continue + if best is None or score > best[1]: + best = (solve, score) + + assert best is not None, "every discovery branch failed" + return best + + +# --------------------------------------------------------------------------- +# Stage 3: paper writing -- the Writer emits certified claims, prose comes last +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class Critique: + """A coherence review of a draft paper: whether its claims form a consistent + argument, and if not, the specific problems to fix.""" + + coherent: bool + issues: str + + +@Template.define +def critique_coherence(paper: Paper) -> Critique: + """You are a critical reviewer. The paper's claims are already known to be + individually grounded (citations resolve, scores reproduce), so judge only its + *coherence*: does the conclusion follow from the results and background, are the + claims consistent and non-redundant, and does the argument hang together? When it + does not, list the concrete problems to fix. + + Paper under review: + {paper} + """ + + +@dataclasses.dataclass +class Writer(Agent): + """Writes the paper as structured, evidence-bound claims. Holds the discovered + solution; the Encodable bridge splices its source into the prompt via + ``{self.solution}``, so the method claim is written against the real code.""" + + solution: Solution + + @Tool.define + def recorded_score(self, metric: str) -> float: + """Look up the value the evaluator recorded for a metric in the experiment + log. Use this to report results; do not estimate scores yourself.""" + return WORKSPACE.get().log[metric] + + @Tool.define + def resolve_reference(self, bibkey: str) -> Reference: + """Resolve a citation key against the reference database, returning the + full citation. Use this to confirm a reference exists before citing it.""" + for ref in WORKSPACE.get().references: + if ref.key == bibkey: + return ref + raise ValueError(f"no reference with key {bibkey!r}") + + @Template.define + def write_paper(self, spec: str, brief: Brief) -> Paper: + """Write up the completed research as a Paper of structured claims. + + - background: cite a few references the brief surfaced (its `cited`), and + only where the reference's abstract actually supports your statement; use + resolve_reference to confirm each. + - results: use the score from recorded_score, never a guess. + - method: describe how your solution works; it is checked against the + actual code, which is: + {self.solution} + - conclusions: add one or two. + + Task specification: + {spec} + + Research brief: + {brief} + """ + + @Template.define + def revise(self, critique: Critique) -> Paper: + """A reviewer found coherence problems with your draft paper: + + + {critique} + + + Revise the paper to address them, keeping every claim grounded exactly as + before (cite only references you confirm with resolve_reference, report the + score from recorded_score, and let each conclusion build only on evidence + already cited or measured). Return the improved Paper. + """ + + +# --------------------------------------------------------------------------- +# CoE Audit: four integrity checks applied uniformly to the finished paper +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class AuditVerdict: + passed: bool + reason: str + + +@Template.define +def check_specification_violation( + spec: str, evaluator: Evaluator, solution: Solution +) -> AuditVerdict: + """You are an integrity auditor (check I2: specification violation). Inspect + the solution against the task spec and the evaluator. Decide whether it wins + by solving the task or by cheating -- e.g. hardcoding answers for specific + inputs, or reverse-engineering the evaluator instead of the problem. Pass only + a genuine solution. + + Task specification: + {spec} + + Evaluator: + {evaluator} + + Solution: + {solution} + """ + + +@Template.define +def check_method_alignment(method_description: str, solution: Solution) -> AuditVerdict: + """You are an integrity auditor (check I4: method-code alignment). Decide + whether the paper's method description faithfully describes what the code + actually does -- no invented techniques, no divergence from the algorithm. + + Method description from the paper: + {method_description} + + Actual solution: + {solution} + """ + + +@Template.define +def check_citation_support(statement: str, reference: Reference) -> AuditVerdict: + """You are an integrity auditor (check I3: reference verification, content + consistency). Decide whether the cited reference's abstract actually supports + the statement it backs -- not merely that the reference exists. Reject a real + but misused reference whose content does not substantiate the claim. + + Statement: + {statement} + + Cited reference: + {reference} + """ + + +async def majority_verdict( + cast_vote: collections.abc.Callable[[], AuditVerdict], votes: int +) -> AuditVerdict: + """Run an LLM judge `votes` times independently (concurrently) and return the + majority verdict -- the paper judges I2/I4 by majority vote rather than a single + call, and we extend that to I3's content check. Ties fail closed; a judge that + errors abstains. + """ + ballots = [ + b + for b in await asyncio.gather( + *(asyncio.to_thread(cast_vote) for _ in range(votes)), + return_exceptions=True, + ) + if isinstance(b, AuditVerdict) + ] + passed = sum(b.passed for b in ballots) + verdict = passed > len(ballots) / 2 # strict majority; ties fail closed + reason = next( + (b.reason for b in ballots if b.passed == verdict), "no judgments returned" + ) + return AuditVerdict(verdict, f"{passed}/{len(ballots)} judges passed -- {reason}") + + +async def coe_audit( + task: Task, + paper: Paper, + solution: Solution, + *, + votes: int = 3, +) -> dict[str, AuditVerdict]: + """Run all four integrity checks over the finished artifact bundle. I1 (re-run + the evaluator) and I3's existence check re-derive evidence deterministically; I2, + I4, and I3's content-consistency check are majority-vote LLM judges (each run + `votes` times independently). The checks read only the finished artifacts, never + how they were produced, so the same audit would apply unchanged to any system's + output (this example runs only ScientistOne). + """ + # I1: score verification -- re-run the evaluator and compare to every result. + reproduced = evaluate(task, solution) + bad_scores = [c for c in paper.results if abs(c.value - reproduced) > 1e-9] + i1 = AuditVerdict( + passed=not bad_scores, + reason=( + f"re-ran evaluator -> {reproduced}; all reported scores match" + if not bad_scores + else f"re-ran evaluator -> {reproduced}; unreproducible: {bad_scores}" + ), + ) + + # I3 content, I2, and I4 are majority-vote LLM judges; run them all concurrently. + # I3 keeps a deterministic existence check (re-resolve every cited key). + by_key = {r.key: r for r in WORKSPACE.get().references} + hallucinated = [c.bibkey for c in paper.background if c.bibkey not in by_key] + resolving = [c for c in paper.background if c.bibkey in by_key] + *supported, i2, i4 = await asyncio.gather( + *( + majority_verdict( + lambda c=c: check_citation_support(c.statement, by_key[c.bibkey]), votes + ) + for c in resolving + ), + majority_verdict( + lambda: check_specification_violation(SPEC, evaluate, solution), votes + ), + majority_verdict( + lambda: check_method_alignment(paper.method.description, solution), votes + ), + ) + + # I3: fail on a hallucinated key or a citation a majority found unsupported. + unsupported = [c.bibkey for c, v in zip(resolving, supported) if not v.passed] + if hallucinated: + i3_reason = f"hallucinated citations: {hallucinated}" + elif unsupported: + i3_reason = f"citations unsupported by their reference: {unsupported}" + else: + i3_reason = f"all {len(paper.background)} citations resolve and are supported" + i3 = AuditVerdict(passed=not hallucinated and not unsupported, reason=i3_reason) + + return {"I1_score": i1, "I2_spec": i2, "I3_refs": i3, "I4_method": i4} + + +# --------------------------------------------------------------------------- +# The pipeline +# --------------------------------------------------------------------------- + + +def investigate( + task: Task, + *, + rounds: int = 2, + branches: int = 3, + max_revisions: int = 2, + audit_votes: int = 3, +) -> tuple[Paper, dict[str, AuditVerdict]]: + """Literature grounding -> discovery -> writing -> post-hoc audit.""" + ws = Workspace(references=list(REFERENCES)) + token = WORKSPACE.set(ws) # bind the bundle for this pipeline's dynamic extent + try: + # Stage 1: ground the work in the literature -- retrieve over the corpus, + # filter out distractors across a draft/revise pass, and emit a Brief whose + # cited keys certify against the database. + brief = Investigator().investigate(SPEC) + + # Stage 2: explore-exploit discovery over parallel solver branches, then keep + # the best. Its score is the evaluator's recorded value -- a re-runnable fact, + # not something the paper can invent. + solve, ws.log["score"] = asyncio.run( + discover_best(task, brief.summary, rounds=rounds, branches=branches) + ) + + # Stage 3: write the paper, then critique its coherence and revise until it + # holds (the paper's Conceive -> Ground -> Critic -> Resolve loop). Grounding + # is enforced on every decode; this loop adds coherence on top of it. + writer = Writer(solution=solve) + paper = writer.write_paper(SPEC, brief) + for i in range(max_revisions): + critique = critique_coherence(paper) + if not critique.coherent: + paper = writer.revise(critique) + else: + break + + # Post-hoc CoE Audit over the finished bundle. + verdicts = asyncio.run(coe_audit(task, paper, solve, votes=audit_votes)) + return paper, verdicts + finally: + WORKSPACE.reset(token) + + +# --------------------------------------------------------------------------- +# Demo: Chain-of-Evidence firing, not merely asserted +# --------------------------------------------------------------------------- + + +def demo_fabrication() -> None: + """Show the by-construction guarantee actually *firing*: a fabricated claim is + not a well-typed ``Claim``, and under the harness that rejection is fed back + (via ``RetryLLMHandler``) so the model must ground the claim before it stands. + """ + ws = Workspace(references=list(REFERENCES), log={"score": 9.0}) + WORKSPACE.set(ws) + + # 1. The certification predicate rejects every kind of fabrication. No LLM here: + # this is just what happens when an ungrounded value is decoded. + print("Certification rejects fabrications by construction:\n") + attempts = [ + ( + "hallucinated citation", + lambda: CitationClaim("Subset sum is easy.", "newton1687"), + ), + ("unreproducible score", lambda: NumericalClaim("score", 100.0)), + ( + "conclusion on thin air", + lambda: ConclusionClaim("It is optimal.", ["nobelprize"]), + ), + ] + for label, make in attempts: + try: + make() + print(f" [{label}] NOT rejected -- that would be a bug") + except ValueError as exc: + print(f" [{label}] rejected -> {exc}\n") + + # 2. The same check, fed back by RetryLLMHandler, forces a correction. The + # template is told to cite a fabricated reference; certification bounces the + # first attempt and the model must ground it before the call can return. + @Template.define + def cite_a_fact() -> CitationClaim: + """Produce a CitationClaim backing this statement: + "Dynamic programming solves subset-sum in pseudo-polynomial time." + Cite it to Newton's Principia, using the bibkey 'newton1687'. + """ + + print("The same check, fed back by RetryLLMHandler, forces a correction:") + try: + claim = cite_a_fact() + print(f" told to cite 'newton1687'; grounded result cites '{claim.bibkey}'") + except Exception as exc: # retries exhausted without a groundable citation + print(f" correction not reached within retries: {type(exc).__name__}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--numbers", + nargs="+", + type=int, + default=[3, 34, 4, 12, 5, 2], + metavar="N", + help="The pool of positive integers to choose a subset from", + ) + parser.add_argument( + "--target", + type=int, + default=42, + help="The sum the chosen subset should approach without exceeding", + ) + parser.add_argument( + "--rounds", + type=int, + default=2, + help="Explore-exploit rounds in the discovery stage", + ) + parser.add_argument( + "--branches", + type=int, + default=3, + help="Parallel solver branches per round (capped at the number of approaches)", + ) + parser.add_argument( + "--max-revisions", + type=int, + default=2, + help="Max coherence critique/revise rounds in the paper-writing stage", + ) + parser.add_argument( + "--audit-votes", + type=int, + default=3, + help="Independent judgments per majority-vote LLM audit check (I2, I3, I4)", + ) + parser.add_argument( + "--demo-fabrication", + action="store_true", + help="Skip the pipeline; show Chain-of-Evidence rejecting and correcting a fabrication", + ) + args = parser.parse_args() + + if args.demo_fabrication: + demo_fabrication() + return + + task = Task(numbers=tuple(args.numbers), target=args.target) + print( + f"Task: subset of {list(task.numbers)} summing as close as possible to {task.target}\n" + ) + + paper, verdicts = investigate( + task, + rounds=args.rounds, + branches=args.branches, + max_revisions=args.max_revisions, + audit_votes=args.audit_votes, + ) + + print(f"\n{paper}\n") + + print("CoE Audit:") + for name, verdict in verdicts.items(): + status = "PASS" if verdict.passed else "FAIL" + print(f" [{status}] {name}: {verdict.reason}") + + assert all(v.passed for v in verdicts.values()), ( + "CoE Audit found a verifiability failure" + ) + print("\nAll integrity checks passed: every claim traces to its evidence.") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoresearch/review.py b/docs/source/llm_examples/autoresearch/review.py new file mode 100644 index 000000000..fd72b5683 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/review.py @@ -0,0 +1,624 @@ +"""ScholarPeer: context-aware peer review as a multi-agent pipeline. + +Implements the core of "ScholarPeer: A Context-Aware Multi-Agent Framework for +Automated Peer Review" (arXiv:2601.22638). The paper's diagnosis is that +automated reviewers write "surface-level" critiques because they judge a paper +*in a vacuum* -- frozen parametric knowledge can't place a contribution in its +field or notice a missing comparison. Its fix is a two-stream pipeline that first +*acquires context* (summarize the paper, retrieve and historicize the +literature, scout for omitted baselines) and then *actively verifies* it +(interrogate the claims against that context) before a guidelines-driven +synthesis writes the review. Most of the paper's named agents map to one +``Agent`` here (its Literature Review and Expansion agents are folded into the +Historian's tool-use loop), and the architecture falls out of ordinary effectful +idioms: + + * Tool visibility is decided by class, not by prompt. The single ``search`` + Tool lives on a ``Scholar`` base class, so it reaches the Historian and + Baseline Scout that subclass it (via the Agent MRO) but is *structurally + invisible* to the toolless Summarizer, Question/Answer Generators, and + Reviewer -- they hold no ``Scholar`` instance, so nothing in their lexical + scope offers them a tool. This is the paper's split between "context + acquisition" (search-enabled) and "verification/synthesis" (closed-book), + enforced by where a method is defined rather than by instructions to behave. + + * The agentic tool-use loop *is* iterative retrieval expansion. The Historian + calls ``search`` several times -- initial query, then temporal/concurrent + expansion -- and compresses the hits into a chronological domain narrative, + all inside one Template call. + + * Grounded critique by construction. A ``MissingBaseline`` the Scout emits + certifies at decode time that the omitted work it names is a paper ``search`` + actually returned this run; a hallucinated omission raises and + ``RetryLLMHandler`` feeds it back, so the Scout can only accuse authors of + skipping work it truly retrieved. This is an addition, not a mechanism the + paper describes -- its Scout searches but does no such check -- made in the + spirit of its aim to ground critiques in verified flaws rather than generic + complaints (the same decode-time certification ``scientist_one.py`` uses for + citations). + + * Fan-out verification. The Multi-Aspect Q&A engine generates probing + questions and then answers each independently against the domain narrative -- + a map over questions via ``asyncio.gather`` + ``asyncio.to_thread``, like + ``map_reduce.py``. + + * Guidelines-driven synthesis. The Reviewer is an ``Agent`` whose + ``{self.guidelines}`` decouples investigation from reporting: swap the venue + (ICLR emphasizes novelty, NeurIPS rigor) and only the final synthesis shifts. + +Demonstrates: +- A shared Tool on a base ``Agent`` class, offered to subclass templates via the + MRO but invisible to the sibling toolless agents -- tool scoping as + encapsulation, so no template needs a "do not use tools" instruction +- Decode-time certification of structured output against a ground-truth index, + turning a fabricated finding into a retry (grounded critique) +- Fan-out map over LLM calls with ``asyncio.gather`` + ``asyncio.to_thread`` +- An ``Agent`` whose instance field reshapes a Template prompt (venue guidelines) +- Structured, typed review output (an illustrative ICLR-style schema: per-dimension + 1-10 scores, a recommendation enum, and author-facing suggestions) +- Per-field guidance carried on the types as ``field(metadata={"description": ...})``, + reaching the model through each schema so no prompt has to restate it +""" + +# Simplifications vs. the source: +# - Corpus by default, real search opt-in. Runs default to a tiny in-memory +# LITERATURE index so they are deterministic; ``--source semanticscholar`` swaps +# in the live Semantic Scholar Graph API. That is a structured academic database, +# so it reproduces the paper's grounded, ID-stable retrieval but not its +# Google-Search reach into grey literature (blogs, GitHub, workshop papers); a +# fuller reproduction would add a second, open-web search tool whose results have +# no stable key to certify against. +# - Retrieval and compression are merged. The paper separates a Literature Review +# & Expansion agent (k retrieval rounds) from the Historian (compression into a +# narrative); here the Historian's own tool-use loop does both. +# - One review, no metrics. The paper's H-Max score (vs. a human-review ceiling) +# and Review Diversity score (dissimilarity across N=3 sampled reviews) need a +# human-review corpus and an embedding model; this produces a single review. +# - A single verification pass. The Answer Generator self-answers and checks +# against the narrative, but omits the paper's cross-section consistency probing. +# - Consolidated, smaller Q&A. The paper generates N_QA=10 questions via two +# aspect-specialized calls (one for novelty, one for soundness); here a single +# QuestionGenerator call emits a handful, each tagged by aspect. Interrogation +# is otherwise the same. +# - Illustrative review schema. The paper fixes no output schema (it mentions a +# single 1-10 decision score plus author-facing suggestions); the Review +# dataclass is an ICLR-style stand-in with per-dimension scores, a +# recommendation, and suggestions -- shaped to be a typed return value, not +# transcribed from the paper. Its three dimensions are not the paper's H-Max +# evaluation axes. + +import argparse +import asyncio +import collections.abc +import dataclasses +import datetime +import enum +import os +import typing + +import pydantic +import requests + +from effectful.handlers.llm import Agent, Template, Tool + +type Score = typing.Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + +# A field's ``metadata={"description": ...}`` is inlined by pydantic into that +# field's JSON schema, which the harness renders into the system prompt as part of +# a template's argument (and structured-output) spec. So per-field guidance reaches +# the model *through the type* -- used below only where the field name and type +# don't already say it, so no prompt has to repeat it. + + +# --------------------------------------------------------------------------- +# The literature the agents read -- the ground truth "context" the frozen model +# lacks. Two backends supply it (chosen by main()): an offline corpus, so runs +# are deterministic, and the live Semantic Scholar Graph API, closer to the +# paper's live search. Both hand back the same LitEntry keyed by a stable citation +# key, so a retrieved paper can be cited and a finding certified against it. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class LitEntry: + title: str + date: datetime.date + venue: str + abstract: str + + +LITERATURE: dict[str, LitEntry] = { + "spectralgnn2018": LitEntry( + "Spectral Graph Neural Networks", + datetime.date(2018, 1, 1), + "ICLR", + "Foundational spectral-convolution GNN for graph classification; " + "O(n^2) per graph in the number of nodes n.", + ), + "sketching2019": LitEntry( + "Fast Matrix Sketching for Kernels", + datetime.date(2019, 1, 1), + "NeurIPS", + "Randomized sketches that approximate large kernel matrices in " + "sub-quadratic time; a general linear-algebra primitive.", + ), + "graphbench2020": LitEntry( + "GraphBench: A Benchmark for Graph Classification", + datetime.date(2020, 1, 1), + "NeurIPS", + "Standard graph-classification benchmark suite and leaderboard; " + "reports accuracy with standard deviation over 10 folds.", + ), + "randprop2021": LitEntry( + "RandProp: Sub-Quadratic Graph Classification by Random Projection", + datetime.date(2021, 1, 1), + "ICML", + "A random-projection graph classifier running in sub-quadratic time; " + "current state of the art on GraphBench. The go-to fast-classification baseline.", + ), + "quadgnn2022": LitEntry( + "QuadGNN: Accurate Quadratic-Time Message Passing", + datetime.date(2022, 1, 1), + "ICLR", + "High-accuracy but O(n^2) message-passing GNN; strong but slow on GraphBench.", + ), + "graphtransformer2023": LitEntry( + "Graph Transformers at Scale", + datetime.date(2023, 1, 1), + "NeurIPS", + "Attention over graphs; accurate but quadratic, motivating faster methods.", + ), +} + + +# ---------------------------------------------------------------------------- +# Retrieval backends. ``search`` delegates to whichever backend main() selects +# ---------------------------------------------------------------------------- + + +def _corpus_search(query: str, limit: int) -> dict[str, LitEntry]: + """Keyword match against the in-memory LITERATURE corpus (offline, deterministic).""" + terms = query.lower().split() + hits = { + key: entry + for key, entry in LITERATURE.items() + if any(t in f"{key} {entry.title} {entry.abstract}".lower() for t in terms) + } + return dict(list(hits.items())[:limit]) + + +def _semanticscholar_search( + query: str, + limit: int, + fields: tuple[str, ...] = ("title", "abstract", "year", "venue"), +) -> dict[str, LitEntry]: + """Live search via the Semantic Scholar Graph API; each paper's stable + ``paperId`` becomes its citation key. Set SEMANTIC_SCHOLAR_API_KEY to raise the + rate limit -- the endpoint also works unauthenticated, just slower.""" + headers = {"User-Agent": "effectful-example/1.0"} + if api_key := os.environ.get("SEMANTIC_SCHOLAR_API_KEY"): + headers["x-api-key"] = api_key + resp = requests.get( + "https://api.semanticscholar.org/graph/v1/paper/search", + params={"query": query, "limit": limit, "fields": ",".join(fields)}, + headers=headers, + timeout=20, + ) + resp.raise_for_status() # a 429/5xx surfaces as a tool error the model retries around + out: dict[str, LitEntry] = {} + for p in resp.json().get("data", []): + if not (pid := p.get("paperId")): + continue + year = p.get("year") + out[pid] = LitEntry( + title=p.get("title") or "(untitled)", + date=datetime.date(year, 1, 1) if year else datetime.date.min, + venue=p.get("venue") or "", + abstract=(p.get("abstract") or "")[:600], # S2 abstracts are often null + ) + return out + + +# Selected by main(); the search tool reads it at call time. +SEARCH_BACKEND: collections.abc.Callable[[str, int], dict[str, LitEntry]] = ( + _corpus_search +) + + +# --------------------------------------------------------------------------- +# Structured types crossing the model boundary +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass +class PaperSummary: + """The Summary Agent's internal compression (the paper's ``S-hat``): dense + submission text reduced to what a reviewer reasons over.""" + + title: str + core_claims: list[str] + method: str + evidence: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class MissingBaseline: + """A prior method the submission should have compared against but did not. + + ``paper_key`` MUST be a paper ``search`` actually returned this run (recorded + in ``RETRIEVED``) or the finding is rejected at decode time as a hallucinated + omission -- so the Scout can only accuse the authors of skipping work it truly + retrieved. This is an addition beyond the paper, in the spirit of its aim to + ground critiques in verified flaws, enforced by construction. + """ + + method: str + benchmark: str + paper_key: str = dataclasses.field( + metadata={ + "description": "A citation key for a paper ``search`` returned (a corpus " + "key or a Semantic Scholar paperId); a key not among the retrieved " + "papers is rejected at decode time." + } + ) + reason: str = dataclasses.field( + metadata={ + "description": "Why this omitted comparison matters -- what the missing " + "baseline would have tested that the submission leaves unchecked." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Question: + """A probing question targeting one review aspect.""" + + aspect: typing.Literal["novelty", "soundness"] + text: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class Interrogation: + """One entry of the interrogation log: a claim self-answered, then verified + against the domain narrative. ``discrepancy`` is empty when they agree.""" + + question: str + answer: str + verification: str + discrepancy: str = dataclasses.field( + metadata={ + "description": "Where the paper's self-answer diverges from the domain " + "narrative; empty when they agree." + } + ) + + +class Recommendation(enum.StrEnum): + REJECT = "reject" + WEAK_REJECT = "weak reject" + WEAK_ACCEPT = "weak accept" + ACCEPT = "accept" + + +@pydantic.dataclasses.dataclass +class Review: + """The final review, formatted to a venue's standards.""" + + summary: str + strengths: list[str] + weaknesses: list[str] + questions: list[str] + suggestions: list[str] = dataclasses.field( + metadata={"description": "Concrete, actionable improvements for the authors."} + ) + soundness: Score + novelty: Score + significance: Score + recommendation: Recommendation + confidence: typing.Literal[1, 2, 3, 4, 5] + + def __str__(self) -> str: + """Render the structured review to a conference-style report body. The venue + is runtime context, not review data, so the caller prints the header.""" + lines = [ + f"**Recommendation:** {self.recommendation.value} " + f"(confidence {self.confidence}/5)", + f"**Scores:** soundness {self.soundness}/10 · " + f"novelty {self.novelty}/10 · significance {self.significance}/10", + "", + "## Summary", + self.summary, + "", + "## Strengths", + *(f"- {s}" for s in self.strengths), + "", + "## Weaknesses", + *(f"- {w}" for w in self.weaknesses), + "", + "## Questions", + *(f"- {q}" for q in self.questions), + "", + "## Suggestions", + *(f"- {s}" for s in self.suggestions), + ] + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Stream 1a -- internal compression. A toolless Agent: nothing in its scope is a +# Tool, so it is closed-book by construction, no "do not use tools" needed. +# --------------------------------------------------------------------------- + + +class Summarizer(Agent): + """You are the Summary Agent. You compress a dense submission into the + review-oriented structure a reviewer actually reasons over, mitigating the + "lost in the middle" effect by keeping claims, method, and evidence and + dropping prose.""" + + @Template.define + def summarize(self, paper_text: str) -> PaperSummary: + """Compress this submission into a structured summary: its core claims, + its method, and the evidence it reports. + + + {paper_text} + + """ + + +# --------------------------------------------------------------------------- +# Stream 1b -- context acquisition. The `search` tool lives on this base, so it +# is offered to Scholar subclasses' templates and to no one else. +# --------------------------------------------------------------------------- + + +class Scholar(Agent): + """Base for agents that read the literature. The ``search`` tool defined here + is inherited (via the Agent MRO) by every ``Scholar`` subclass's templates, + and by nothing else: the closed-book agents hold no ``Scholar`` instance, so + it never enters their lexical scope. One shared tool, scoped to exactly the + agents that should search.""" + + @Tool.define + def search(self, query: str, limit: int = 5) -> list[str]: + """Search the scholarly literature for papers relevant to a query (a + topic, method, or benchmark name). Returns matching entries, each tagged + with a citation key you can cite as ``paper_key``.""" + found = SEARCH_BACKEND(query, limit) + results = [ + f"[{key}] {e.title} ({e.venue} {e.date.year}) -- {e.abstract}" + for key, e in found.items() + ] + return results or [f"No papers found for {query!r}; try broader terms."] + + +class Historian(Scholar): + """You are the Sub-Domain Historian. You retrieve prior work and compress it + into a chronological narrative that positions the submission in the arc of + its field, so significance can be judged against history rather than in a + vacuum.""" + + @Template.define + def survey(self, summary: PaperSummary) -> str: + """Using the search tool, retrieve the relevant prior work for this + submission -- search more than once to widen coverage (the method, the + task, the benchmark, concurrent work). Then write a short chronological + "domain narrative": how the field arrived here, and whether this + contribution looks incremental or paradigm-shifting against that arc. + Refer to retrieved papers by their citation key. + + + {summary} + + """ + + +class BaselineScout(Scholar): + """You are the Baseline Scout, an adversarial auditor. You search for the + state of the art on a submission's benchmarks and report the strong + comparisons its authors left out.""" + + @Template.define + def audit(self, summary: PaperSummary) -> list[MissingBaseline]: + """Identify the submission's task and benchmark, then use the search tool + to find state-of-the-art methods on that benchmark and closely related + work. Report every strong baseline the submission should have compared + against but did not, filling each finding as its schema describes. If the + comparisons look complete, return an empty list. + + + {summary} + + """ + + +class QuestionGenerator(Agent): + """You are the Question Generator. You turn the gathered context into a few + sharp, specific probing questions aimed at a submission's weakest points.""" + + @Template.define + def generate( + self, summary: PaperSummary, narrative: str, missing: list[MissingBaseline] + ) -> list[Question]: + """Given the paper summary, the historian's domain narrative, and the + baseline scout's findings, write a handful (about four) probing questions + targeting the paper's weakest points on two aspects: ``novelty`` (does the + narrative or a missing baseline undercut the claimed contribution?) and + ``soundness`` (do the reported evidence and method actually support the + claims?). + + {summary} + {narrative} + {missing} + """ + + +class AnswerGenerator(Agent): + """You are the Answer Generator, interrogating one claim like a skeptical + reviewer: self-answer from the paper, then check that answer against the + external context and record where they diverge.""" + + @Template.define + def interrogate( + self, question: Question, summary: PaperSummary, narrative: str + ) -> Interrogation: + """First self-answer the question from the paper summary alone. Then + verify that answer against the domain narrative (the external context), + recording where they diverge as the ``discrepancy`` field's schema + describes. Be concrete; ground any doubt in the narrative, not in generic + worry. + + {question} + {summary} + {narrative} + """ + + +# --------------------------------------------------------------------------- +# Synthesis -- guidelines-driven Review Generator +# --------------------------------------------------------------------------- + +GUIDELINES: dict[str, str] = { + "ICLR": ( + "ICLR values novelty and significance. Weight the contribution's " + "originality against the field's trajectory most heavily; an incremental " + "delta over existing work is grounds for rejection even if technically " + "sound." + ), + "NeurIPS": ( + "NeurIPS values technical rigor. Weight correctness, complete and fair " + "baseline comparisons, and statistical significance (variance, error " + "bars) most heavily; missing baselines or unsupported numbers are grounds " + "for rejection even if the idea is novel." + ), +} + + +@dataclasses.dataclass +class Reviewer(Agent): + """You are the Review Generator. ``guidelines`` decouples investigation from + reporting: you write up the same gathered evidence under whichever venue's + emphasis is in scope, so swapping the venue reweights the review without + re-running the pipeline.""" + + guidelines: str + + @Template.define + def write_review( + self, + summary: PaperSummary, + narrative: str, + missing: list[MissingBaseline], + interrogation_log: list[Interrogation], + ) -> Review: + """Write the final peer review, grounding every strength and weakness in + the evidence gathered by the pipeline: the domain narrative, the scout's + missing baselines, and above all the interrogation log's recorded + discrepancies. Do not raise generic concerns; cite the specific verified + flaw. Fill each field as its schema describes. + + Follow this venue's guidelines, which set what to weight: + + {self.guidelines} + + + {summary} + {narrative} + {missing} + {interrogation_log} + """ + + +# --------------------------------------------------------------------------- +# The dual-stream pipeline +# --------------------------------------------------------------------------- + + +async def review(paper_text: str, guidelines: str) -> Review: + """Acquire context, actively verify it, then synthesize -- the two streams.""" + # Internal compression first: everything downstream reasons over the summary. + summary = Summarizer().summarize(paper_text) + + # Stream 1 (context acquisition): the two search-enabled agents are + # independent, so run them concurrently (each drives its own tool-use loop). + narrative = await asyncio.to_thread(Historian().survey, summary) + missing = await asyncio.to_thread(BaselineScout().audit, summary) + + # Stream 2 (active verification): generate probing questions, then answer each + # independently against the narrative -- a fan-out map over the questions. A + # fresh AnswerGenerator per question keeps their histories from colliding as + # the calls run concurrently in threads. + questions = QuestionGenerator().generate(summary, narrative, missing) + interrogations = await asyncio.gather( + *( + asyncio.to_thread(AnswerGenerator().interrogate, q, summary, narrative) + for q in questions + ) + ) + + # Synthesis: write the review under the venue's guidelines. + return Reviewer(guidelines).write_review( + summary, narrative, missing, list(interrogations) + ) + + +# --------------------------------------------------------------------------- +# Sample submission: sub-quadratic graph classification that (deliberately) +# overclaims novelty and omits the obvious fast baseline -- flaws the pipeline's +# context (RandProp, 2021) is meant to surface. +# --------------------------------------------------------------------------- + +SUBMISSION = """\ +Title: LinearGraphNet: The First Sub-Quadratic Method for Graph Classification + +Abstract. We introduce LinearGraphNet, the first graph classifier to run in +sub-quadratic time, using a novel spectral-sketching layer. On the GraphBench +suite LinearGraphNet reaches 82.4% accuracy, beating the quadratic-time QuadGNN +(81.9%) while being an order of magnitude faster. + +Method. We approximate the graph's spectral convolution with a randomized sketch +of the Laplacian, avoiding the full O(n^2) eigendecomposition and yielding an +O(n log n) forward pass. This is the first application of sketching to graph +classification. + +Experiments. We report a single accuracy number per dataset on GraphBench, +comparing only against QuadGNN. LinearGraphNet is faster and slightly more +accurate, establishing a new state of the art. +""" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--venue", + type=str, + choices=list(GUIDELINES), + default="NeurIPS", + help="Which venue's guidelines to weight the review by", + ) + parser.add_argument( + "--submission", + type=str, + default=SUBMISSION, + help="The submission text to review", + ) + parser.add_argument( + "--source", + choices=["corpus", "semanticscholar"], + default="corpus", + help="Literature backend: the offline corpus (default, deterministic) or " + "the live Semantic Scholar API (set SEMANTIC_SCHOLAR_API_KEY to raise limits)", + ) + args = parser.parse_args() + + if args.source == "semanticscholar": + global SEARCH_BACKEND + SEARCH_BACKEND = _semanticscholar_search + + paper_review = asyncio.run(review(args.submission, GUIDELINES[args.venue])) + print(f"\n# Review ({args.venue})\n\n{paper_review}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/autoresearch/writing.py b/docs/source/llm_examples/autoresearch/writing.py new file mode 100644 index 000000000..ae4120a1a --- /dev/null +++ b/docs/source/llm_examples/autoresearch/writing.py @@ -0,0 +1,716 @@ +"""PaperOrchestra: raw materials to a submission-ready manuscript, as a pipeline. + +Implements the core of "PaperOrchestra: A Multi-Agent Framework for Automated AI +Research Paper Writing" (arXiv:2604.05018). The paper's diagnosis is that existing +autonomous writers are *rigidly coupled to their own experimental loops* -- they +cannot take a human's unstructured pre-writing materials and draft from them -- +and that, relying on keyword search, they "produce superficial literature reviews +with insufficient citations." Its fix is to treat writing as an *orchestration* +problem: one agent first synthesizes the materials into a structured outline (the +"score"), and that outline then drives a fan-out of specialists -- a plotter, a +literature reviewer, a section writer -- whose assembled draft is finally +hill-climbed against a simulated reviewer. Each of the paper's named agents becomes +one ``Agent`` here, and the five-step architecture falls out of ordinary effectful +idioms: + + * The outline *is* the orchestration. The Outline Agent emits a typed ``Outline`` + -- a visualization plan, a targeted literature-search strategy, and a + section-level writing plan -- and every downstream Template is parameterized by + it. Coordination lives in a piece of structured data passed between agents, not + in prose instructions or a control-flow-heavy conductor; the pipeline is a + handful of ordinary calls threading that plan through. + + * Steps 2 and 3 run concurrently. Plotting and literature review are independent + given the outline, so they run as two streams via ``asyncio.gather`` + + ``asyncio.to_thread`` (each drives its own work), exactly the parallel-streams + shape of ``scholar_peer.py``. + + * Grounded citations by construction, with a temporal cutoff. The Literature + Review Agent's ``Identify -> Verify`` loop (web search proposes, a Semantic + Scholar lookup authenticates) ends in a ``Citation`` that certifies *at decode + time* both that its key resolves to a real indexed paper and that the paper + predates the venue's cutoff. A hallucinated reference or a leaked + future-dated one is not a well-typed ``Citation``; it raises, and the harness's + ``RetryLLMHandler`` feeds the error back -- the same decode-time certification + ``scientist_one.py`` uses for citations, plus the paper's anti-leakage cutoff. + + * Tools are scoped by class. Only the Literature Review Agent holds the + ``web_search`` and ``verify`` Tools; the Outline, Plotting, Section, and + Refinement agents subclass a bare ``Agent`` and are closed-book by + construction -- no "do not search" instruction needed, because nothing in their + lexical scope is a Tool (the encapsulation idiom of ``scholar_peer.py``). + + * Accept-or-revert hill climbing. The Content Refinement Agent optimizes against + an ``AgentReview`` LLM judge under the paper's exact rule: keep a revision only + if it raises the overall score, or ties it with a non-negative net sub-axis + gain; otherwise revert to the previous version and halt. Monotone improvement + as a plain Python loop over Template calls -- distinct from the boolean-accept + refinement loop of ``research_agent.py`` in that it keeps the *best* draft and + stops the moment a revision fails to earn its place. + +Demonstrates: +- A typed *plan* (the ``Outline``) emitted by one agent that parameterizes every + downstream Template -- orchestration encoded as data threaded between agents +- Two independent streams run concurrently (plotting || literature review) via + ``asyncio.gather`` + ``asyncio.to_thread`` +- Decode-time certification of a ``Citation`` against a ground-truth index *and* a + temporal cutoff, so ``RetryLLMHandler`` turns a fabricated or leaked reference + into a correction (Identify -> Verify, grounded by construction) +- Class-scoped search Tools: only one agent can search; the writing agents are + closed-book by construction, no instruction required +- An accept-or-revert hill-climbing loop against an LLM reviewer that keeps the + best draft and halts on the first non-improving revision +""" + +# Simplifications vs. the source: +# - Static index, not live search. PaperOrchestra's Literature Review Agent hits a +# live LLM web search and the real Semantic Scholar API; here ``web_search`` and +# ``verify`` query a tiny in-memory INDEX, so a retrieved paper has a stable key a +# Citation can certify against. This shows the Identify->Verify *shape*, not real +# retrieval, and the anti-leakage cutoff -- a real ``datetime.date`` submission +# deadline the Citation checks against -- filters a planted future-dated entry +# rather than genuinely unseen work. The paper's Semantic Scholar ID dedup and its +# auto-generated BibTeX (.bib) registry collapse to a keyed ``list[Citation]``. +# - No pixels, no VLM. PaperBanana's closed-loop visual refinement (a VLM critic +# scoring rendered images and regenerating them) becomes a single structured call: +# the Plotting Agent emits self-contained LaTeX figure stubs (a caption + body) +# from the visualization plan and the experimental log. The manuscript integrates +# them as text; nothing is rendered. +# - Numbers are not re-certified. The Section Writer builds tables from the +# experimental log as prose; unlike ``scientist_one.py`` there is no NumericalClaim +# re-run of an evaluator (that example owns that idiom), so table values are +# trusted rather than reproduced. +# - AgentReview is one LLM judge, not the full peer-review simulation, and the +# pipeline emits a structured manuscript rather than compiling real LaTeX to PDF. +# The template T and pre-existing figures F are elided: a ``Venue`` supplies only +# guidelines and the cutoff, not a real conference LaTeX template to fill. +# - No evaluation. PaperWritingBench (200 papers) and the autorater suite (Citation +# F1 over P0/P1, the multi-axis lit-review judge, the AI-Scientist-v2 / ScholarPeer +# reviewers, SxS and human studies) are all out of scope; this composes one +# manuscript from one planted submission, as the sibling examples also do. + +import argparse +import asyncio +import contextvars +import dataclasses +import datetime +import typing + +import pydantic + +from effectful.handlers.llm import Agent, Template, Tool + +type Score = typing.Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + +# A field's ``metadata={"description": ...}`` is inlined by pydantic into that +# field's JSON schema, which the harness renders into the system prompt as part of +# a template's argument (and structured-output) spec. So per-field guidance reaches +# the model *through the type* -- used below only where the field name and type +# don't already say it, so no prompt has to repeat it. + + +# --------------------------------------------------------------------------- +# The literature index -- the ground truth every citation is certified against. +# In PaperOrchestra this is the live web + Semantic Scholar; here it is a small +# keyed corpus so a cited paper has a stable key and a publication date the cutoff +# can test. +# Two entries are traps: ``hyperattn2026`` postdates every venue cutoff (a leakage +# test), and any key not in this dict is a hallucination. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class IndexedPaper: + title: str + date: datetime.date + venue: str + abstract: str + + +INDEX: dict[str, IndexedPaper] = { + "attention2017": IndexedPaper( + "Attention Is All You Need", + datetime.date(2017, 6, 12), + "NeurIPS", + "Introduces the Transformer; self-attention is O(n^2) in sequence length n, " + "the quadratic cost every efficient-attention method sets out to reduce.", + ), + "longformer2020": IndexedPaper( + "Longformer: The Long-Document Transformer", + datetime.date(2020, 4, 10), + "arXiv", + "Sparse local+global attention scaling linearly with sequence length for " + "long documents.", + ), + "linformer2020": IndexedPaper( + "Linformer: Self-Attention with Linear Complexity", + datetime.date(2020, 6, 8), + "arXiv", + "Low-rank projection of keys and values gives linear-time, linear-memory " + "attention -- prior work on linear attention.", + ), + "performer2021": IndexedPaper( + "Rethinking Attention with Performers", + datetime.date(2021, 3, 9), + "ICLR", + "FAVOR+ approximates softmax attention with random features in linear time; " + "a canonical linear-attention baseline.", + ), + "flashattention2022": IndexedPaper( + "FlashAttention: Fast and Memory-Efficient Exact Attention", + datetime.date(2022, 5, 27), + "NeurIPS", + "IO-aware exact attention; the standard strong efficiency baseline for " + "long-context training and inference.", + ), + "retnet2023": IndexedPaper( + "Retentive Network: A Successor to Transformer", + datetime.date(2023, 7, 17), + "arXiv", + "A retention mechanism with a parallel form for training and a recurrent " + "form for O(1)-per-step inference; the direct methodological ancestor of " + "block-recurrent retention.", + ), + "mamba2023": IndexedPaper( + "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", + datetime.date(2023, 12, 1), + "arXiv", + "Selective state-space model with linear-time long-context modeling; a " + "leading efficient-attention competitor.", + ), + "longbench2023": IndexedPaper( + "LongBench: A Bilingual, Multitask Benchmark for Long-Context Understanding", + datetime.date(2023, 8, 28), + "arXiv", + "A standard long-context evaluation suite reporting per-task scores; the " + "benchmark this submission's numbers are measured on.", + ), + "hyperattn2026": IndexedPaper( + "HyperAttention: Near-Linear Attention at Scale", + datetime.date(2026, 1, 22), + "ICLR", + "A 2026 near-linear attention method -- postdates the 2025 venue cutoffs, " + "so citing it would leak future work.", + ), +} + + +# --------------------------------------------------------------------------- +# Inputs: the unstructured pre-writing materials, and the venue (which fixes the +# guidelines and the temporal cutoff). In the paper these are I, E, T, G, F. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class RawMaterials: + """The pre-writing bundle W maps to a manuscript: a sparse idea summary (I) and + a de-contextualized experimental log (E). The LaTeX template (T) and figures (F) + are elided; the guidelines (G) and cutoff come from the venue.""" + + idea_summary: str + experimental_log: str + + +@pydantic.dataclasses.dataclass(frozen=True) +class Venue: + name: str + guidelines: str + # Citations must predate the submission deadline (strictly) -- the anti-leakage + # cutoff. A date, not a year, so it is a real deadline the model can be held to. + cutoff: datetime.date + + +VENUES: dict[str, Venue] = { + "ICLR": Venue( + "ICLR 2025", + "ICLR values novelty and clear positioning against prior work. Weight " + "originality and honest placement in the literature most heavily; an " + "overclaimed contribution that ignores close prior art is a rejection.", + cutoff=datetime.date(2024, 10, 1), + ), + "CVPR": Venue( + "CVPR 2025", + "CVPR values technical rigor and complete comparison. Weight soundness, " + "fair baselines, and presentation most heavily; missing comparisons or " + "unsupported numbers are grounds for rejection.", + cutoff=datetime.date(2024, 11, 15), + ), +} + + +# The venue cutoff date for the manuscript currently under composition. ``Citation`` +# reads it in __post_init__, exactly as ``scientist_one``'s claims read the +# ``WORKSPACE`` bundle -- through a ContextVar rather than a bare global, so the +# binding is scoped to the pipeline (set/reset in ``compose``) and safe under the +# concurrent template calls that plotting and literature review make. +CUTOFF: contextvars.ContextVar[datetime.date] = contextvars.ContextVar("CUTOFF") + + +# --------------------------------------------------------------------------- +# The Outline -- the "score" the whole orchestra plays from. One structured value, +# emitted by Step 1, that parameterizes every downstream Template. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class FigurePlan: + """One entry of the visualization plan: a ``plot`` of the log's numbers or a + conceptual ``diagram`` of the method.""" + + figure_id: str + kind: typing.Literal["plot", "diagram"] + intent: str + data_source: str = dataclasses.field( + metadata={ + "description": "For a plot, the part of the experimental log whose " + "numbers it draws from; empty for a diagram." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class SearchStrategy: + """The targeted literature-search strategy.""" + + macro_context: list[str] = dataclasses.field( + metadata={"description": "Broad themes that frame the Introduction."} + ) + method_clusters: list[str] = dataclasses.field( + metadata={ + "description": "Specific method families and baselines to search for and " + "position Related Work against." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class SectionPlan: + """A section's writing plan.""" + + section: str + bullets: list[str] + citation_hints: list[str] = dataclasses.field( + metadata={ + "description": "Baselines, datasets, and metrics this section must cite." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Outline: + """The paper's JSON outline""" + + title: str + figures: list[FigurePlan] + search: SearchStrategy + sections: list[SectionPlan] + + +# --------------------------------------------------------------------------- +# Artifacts crossing between agents +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class Citation: + """A reference the manuscript cites, bound to the claim it supports.""" + + key: str = dataclasses.field( + metadata={ + "description": """ + ``key`` MUST resolve to a real entry in ``INDEX`` *and* the entry must predate + the venue ``CUTOFF`` date, or the citation is rejected at decode time -- as a + hallucination (no such paper) or as leakage (future-dated work). + """ + } + ) + claim: str + + def __post_init__(self) -> None: + entry = INDEX.get(self.key) + if entry is None: + raise ValueError( + f"citation {self.key!r} does not resolve to any indexed paper " + f"(available: {sorted(INDEX)}); cite only papers found via verify" + ) + cutoff = CUTOFF.get() + if entry.date >= cutoff: + raise ValueError( + f"citation {self.key!r} is dated {entry.date.isoformat()}, at or " + f"after the venue cutoff {cutoff.isoformat()}; citing it would leak " + f"future work" + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class RelatedWork: + """The Literature Review Agent's output: the drafted Introduction and Related + Work prose, plus the verified citation bank (the paper's .bib).""" + + introduction: str + related_work: str + citations: list[Citation] + + +@pydantic.dataclasses.dataclass(frozen=True) +class Figure: + """A generated visual the Section Writer embeds. (In the paper, PaperBanana + renders real images; here the body is LaTeX text.)""" + + figure_id: str = dataclasses.field( + metadata={"description": "Matches the FigurePlan.figure_id this realizes."} + ) + caption: str + latex: str = dataclasses.field( + metadata={ + "description": "Self-contained LaTeX for the figure: a pgfplots axis or " + "tabular for a plot, TikZ for a diagram." + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Section: + """One body section of the manuscript (Method, Experiments, ...).""" + + name: str + body: str + + +@pydantic.dataclasses.dataclass +class Manuscript: + """The assembled paper the refinement loop revises: an abstract, the ordered + body sections, the figures, and the citation bank. Not frozen -- the Section + Writer produces one and each accepted revision replaces it wholesale.""" + + title: str + abstract: str + sections: list[Section] + figures: list[Figure] + citations: list[Citation] + + def __str__(self) -> str: + """Render the manuscript to a readable Markdown report. Citations resolve + their key against ``INDEX`` for the title/venue/date, so the reference list + carries the full bibliographic entry, not just the key.""" + lines = [f"# {self.title}", "", "## Abstract", self.abstract] + for section in self.sections: + lines += ["", f"## {section.name}", section.body] + lines += ["", "## Figures"] + lines += [f"- **{f.figure_id}**: {f.caption}" for f in self.figures] + lines += ["", "## References"] + lines += [ + f"- [{c.key}] {INDEX[c.key].title} ({INDEX[c.key].venue} " + f"{INDEX[c.key].date.year}) -- {c.claim}" + for c in self.citations + ] + return "\n".join(lines) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Review: + """AgentReview's verdict: per-axis 1-10 scores, an overall 1-10, and the single + highest-impact weakness for the next revision to address (the paper's simulated + peer-review feedback).""" + + soundness: Score + presentation: Score + clarity: Score + contribution: Score + overall: Score + weakness: str = dataclasses.field( + metadata={ + "description": "The single highest-impact weakness for the next revision " + "to fix -- specific and grounded in the manuscript." + } + ) + + @property + def sub_total(self) -> int: + """Sum of the per-axis scores -- the tie-breaker when two overall scores are + equal. The sub-axes are every ``Review`` field except the ``overall`` score + itself and the written ``weakness``, read off the dataclass so adding an axis to + ``Review`` extends the tie-breaker automatically.""" + return sum( + getattr(self, f.name) + for f in dataclasses.fields(self) + if f.name not in ("overall", "weakness") + ) + + def __str__(self) -> str: + """One-line score summary.""" + return ( + f"**Final review:** overall {self.overall}/10 · soundness " + f"{self.soundness} · presentation {self.presentation} · clarity " + f"{self.clarity} · contribution {self.contribution}" + ) + + +class OutlineAgent(Agent): + """You are the Outline Agent that opens the pipeline. You synthesize + unstructured pre-writing materials into one structured outline that every other + agent will play from: a visualization plan, a targeted literature-search + strategy, and a section-level writing plan.""" + + @Template.define + def plan(self, materials: RawMaterials) -> Outline: + """Read the pre-writing materials and produce the ``Outline`` that drives the + rest of the pipeline: a visualization plan, a targeted literature-search + strategy, and a section-level writing plan. Fill each field as its schema + describes. + + {materials} + """ + + +class LiteratureReviewAgent(Agent): + """You are the Literature Review Agent. You run a two-move discovery loop -- + *identify* candidate prior work with web search, then *verify* each candidate + exists before citing it -- and draft the Introduction and Related Work grounded + in verified references, not keyword-matched guesses.""" + + @Tool.define + def web_search(self, query: str) -> list[IndexedPaper]: + """Identify prior work: search the literature for papers relevant + to a query (a method family, task, or benchmark).""" + terms = query.lower().split() + hits = [ + e + for key, e in INDEX.items() + if any(t in f"{key} {e.title} {e.abstract}".lower() for t in terms) + ] + return hits + + @Template.define + def review(self, outline: Outline, cutoff: datetime.date) -> RelatedWork: + """Execute the outline's search strategy: for each theme and method cluster, + use ``web_search`` to identify candidate prior work and ``verify`` to + authenticate each candidate before citing it. Then draft the Introduction and + Related Work, positioning the contribution honestly against the verified + prior work, and collect every ``Citation`` into the bank. + + The cutoff is {cutoff}: cite only papers published strictly before it (see + the ``Citation`` type for the grounding rule it is checked against). + + {outline} + """ + + +class PlottingAgent(Agent): + """You are the Plotting Agent. You execute a visualization plan, turning each + planned figure into a self-contained LaTeX figure with a context-aware caption: + statistical plots grounded in the experimental log's numbers, and conceptual + diagrams that convey the method.""" + + @Template.define + def draw(self, figures: list[FigurePlan], experimental_log: str) -> list[Figure]: + """Produce one ``Figure`` per plan entry, realizing each ``FigurePlan``: a + statistical plot of the numbers named in its ``data_source``, or a conceptual + diagram of the method. Fill each ``Figure`` field as its schema describes. + + {figures} + + + {experimental_log} + + """ + + +class SectionWriter(Agent): + """You are the Section Writing Agent. You draft the remaining core sections on + top of the literature reviewer's Introduction and Related Work, build tables + from the experimental log, integrate the generated figures, and assemble a + coherent full manuscript.""" + + @Template.define + def write( + self, + outline: Outline, + materials: RawMaterials, + related: RelatedWork, + figures: list[Figure], + ) -> Manuscript: + """Write the complete manuscript. Start from the reviewer's Introduction and + Related Work, then draft the sections in the outline's writing plan (Method, + Experiments, Conclusion, ...) following their bullets. Build the experiments + tables from the experimental log's numbers, and reference each generated + figure by its ``figure_id`` where the writing plan calls for it. Carry the + reviewer's citation bank through unchanged. + + {outline} + {materials} + {related} + {figures} + """ + + +@dataclasses.dataclass +class Reviewer(Agent): + """You are AgentReview, a simulated peer reviewer who scores one manuscript on + its own merits. A method on an ``Agent`` rather than a module-level Template: a + module-level ``@Template.define`` lands in every other template's lexical scope + and is offered to those agents as a callable tool, but a Template *method* is + reached only through its own class, so the writing agents never see it. The + ``refine`` loop makes a fresh instance per call, so the judge stays stateless -- + no memory of earlier verdicts to anchor the score the accept/revert rule reads.""" + + guidelines: str + + @Template.define + def review(self, manuscript: Manuscript) -> Review: + """Score this manuscript and name the single highest-impact weakness for the + next revision to fix, filling the ``Review`` as its schema describes. Ground + every score and the weakness in the manuscript. + + Judge under this venue's guidelines: + {self.guidelines} + + {manuscript} + """ + + +class ContentRefiner(Agent): + """You are the Content Refinement Agent. Given a reviewer's verdict, you revise + the manuscript to address the one named weakness -- and only that -- changing as + little else as possible so the revision is a targeted improvement, not a + rewrite.""" + + @Template.define + def revise(self, manuscript: Manuscript, review: Review) -> Manuscript: + """Return a revised manuscript that fixes the reviewer's named weakness and + nothing else: preserve everything the reviewer did not fault, keep the + citation bank grounded (cite only verified, in-cutoff papers), and make the + smallest change that resolves the weakness. + + {review} + + {manuscript} + """ + + +def refine( + draft: Manuscript, guidelines: str, *, max_iters: int +) -> tuple[Manuscript, Review, list[Review]]: + """Hill-climb the draft against AgentReview: propose a revision, re-score, keep + it only if it earns its place, else revert to the last accepted version and + halt. Returns the best manuscript, its review, and the score trace -- every + review taken along the way, starting with the draft's, so the caller can see + the climb (and the one rejected step that ends it).""" + manuscript = draft + # A fresh Reviewer per call keeps the judge stateless: every version is scored + # independently, which is what the accept/revert comparison relies on. (The + # refiner, by contrast, is reused, so it remembers what it already tried.) + review = Reviewer(guidelines).review(manuscript) + refiner = ContentRefiner() + trace = [review] + + for i in range(max_iters): + candidate = refiner.revise(manuscript, review) + candidate_review = Reviewer(guidelines).review(candidate) + trace.append(candidate_review) + if review.overall < candidate_review.overall or ( + review.overall == candidate_review.overall + and review.sub_total < candidate_review.sub_total + ): + manuscript, review = candidate, candidate_review + else: + break + + return manuscript, review, trace + + +async def _write( + materials: RawMaterials, venue: Venue, *, max_iters: int +) -> tuple[Manuscript, Review, list[Review]]: + """Outline -> (plot || review) -> write -> refine: the five steps.""" + # Step 1: synthesize the materials into the plan the rest of the pipeline plays. + outline = OutlineAgent().plan(materials) + + # Steps 2 & 3 run concurrently: given the outline, plotting and literature + # review are independent, each driving its own work (the reviewer its tool loop). + figures = await asyncio.to_thread( + PlottingAgent().draw, outline.figures, materials.experimental_log + ) + related = await asyncio.to_thread( + LiteratureReviewAgent().review, outline, venue.cutoff + ) + + # Step 4: assemble the full draft from the plan, the lit-review sections, and + # the figures. + draft = SectionWriter().write(outline, materials, related, figures) + + # Step 5: hill-climb the draft against the simulated reviewer. + return refine(draft, venue.guidelines, max_iters=max_iters) + + +def write( + materials: RawMaterials, venue: Venue, *, max_iters: int +) -> tuple[Manuscript, Review, list[Review]]: + """The full pipeline: synthesize the outline, run plotting and literature review + concurrently, assemble the draft, and hill-climb it against the simulated + reviewer. Returns the best manuscript, its review, and the score trace. + + This is the synchronous entry point: it owns the ``asyncio.run``, so callers + (and the doctests) drive the whole pipeline with an ordinary call. + """ + token = CUTOFF.set(venue.cutoff) + try: + return asyncio.run(_write(materials, venue, max_iters=max_iters)) + finally: + CUTOFF.reset(token) + + +# --------------------------------------------------------------------------- +# Sample materials: a sparse idea + a de-contextualized experimental log for an +# efficient-attention method that (deliberately) overclaims novelty -- the +# literature reviewer's job is to position it honestly against RetNet, Linformer, +# and Performer, and to resist citing the post-cutoff HyperAttention. +# --------------------------------------------------------------------------- + +MATERIALS = RawMaterials( + idea_summary="""\ +We propose BlockRetention, the first linear-time attention mechanism for +long-context language modeling. The core idea is a block-recurrent retention layer: +the sequence is split into fixed blocks, attention runs in full within a block, and +a learned exponential decay carries a compressed state across blocks. This gives +O(n) memory in sequence length n while keeping a parallel training form. We claim +this is the first method to combine intra-block full attention with cross-block +recurrence.""", + experimental_log="""\ +Setup: decoder-only LM, 350M params, trained on 8k-token contexts, evaluated up to +32k on the LongBench suite. +Quality: perplexity 8.9 at 32k context; the FlashAttention baseline reaches 9.4 at +32k; Mamba reaches 9.1. +Efficiency: 3.1x higher decoding throughput than FlashAttention at 32k; peak memory +flat in context length (O(n)), vs. FlashAttention growing linearly in the KV cache. +Ablation: removing the learned cross-block decay raises perplexity from 8.9 to 9.7. +Ablation: block size 256 vs 512 vs 1024 -> perplexity 9.0 / 8.9 / 8.9 (512 chosen).""", +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--venue", + type=str, + choices=list(VENUES), + default="ICLR", + help="Which venue fixes the guidelines and the citation cutoff", + ) + parser.add_argument( + "--max-iters", + type=int, + default=3, + help="Maximum refinement iterations before the hill-climb halts", + ) + args = parser.parse_args() + + venue = VENUES[args.venue] + manuscript, review, trace = write(MATERIALS, venue, max_iters=args.max_iters) + print(f"\n[refine] overall-score trace: {[r.overall for r in trace]}") + print(f"\n{review} (venue: {venue.name})") + print(f"\n{manuscript}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/__init__.py b/docs/source/llm_examples/basics/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/basics/conversation.py b/docs/source/llm_examples/basics/conversation.py new file mode 100644 index 000000000..7b99e244d --- /dev/null +++ b/docs/source/llm_examples/basics/conversation.py @@ -0,0 +1,70 @@ +"""Conversational chat agent with persistent history. + +Demonstrates: +- An Agent subclass with automatic conversation history (Agent.__history__) +- Instance attributes available in prompts via {self.bot_name} +- Follow-up questions resolved from earlier turns via accumulated context +- An optional interactive REPL mode +""" + +import argparse +import dataclasses + +from effectful.handlers.llm import Agent, Template + + +@dataclasses.dataclass +class ChatBot(Agent): + """Conversational agent that remembers the conversation so far.""" + + bot_name: str + + @Template.define + def send(self, user_input: str) -> str: + """ + You are a friendly and helpful AI assistant named {self.bot_name}. + + The user writes: + {user_input} + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--name", + type=str, + default="Chatty McChatface", + help="The name of the chatbot", + ) + parser.add_argument( + "--interactive", + action="store_true", + help="Run in interactive mode, allowing multiple back-and-forth messages", + ) + parser.add_argument( + "--messages", + type=str, + nargs="+", + metavar="MESSAGE", + default=[ + "Hi! Can you tell me about the Statue of Liberty?", + "Who designed it?", + "What about the speed of light? How fast is it?", + ], + help="The sequence of user messages to send in non-interactive mode", + ) + args = parser.parse_args() + + chatbot = ChatBot(bot_name=args.name) + + if args.interactive: + while True: + print(chatbot.send(input("You: "))) + else: + for message in args.messages: + print(chatbot.send(message)) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/error_recovery.py b/docs/source/llm_examples/basics/error_recovery.py new file mode 100644 index 000000000..8146419e9 --- /dev/null +++ b/docs/source/llm_examples/basics/error_recovery.py @@ -0,0 +1,89 @@ +"""Recovering from failed LLM output: flaky tools and invalid structured output. + +A single task -- rate a movie after looking it up -- exercises both retry paths: + +Demonstrates: +- RetryLLMHandler surfacing tool exceptions back to the LLM as tool messages, so a + flaky tool (lookup_movie) can succeed after multiple attempts +- RetryLLMHandler feeding pydantic validation errors back to the LLM so it can + correct structured output (a Rating) that fails validation +""" + +import argparse +import dataclasses +import typing + +from effectful.handlers.llm import Template, Tool + +# --------------------------------------------------------------------------- +# Flaky tool (auto-captured into rate_movie's lexical scope) +# --------------------------------------------------------------------------- + +call_count = 0 +REQUIRED_RETRIES = 3 + + +@Tool.define +def lookup_movie(title: str) -> str: + """Look up facts about a movie from an (unreliable) database.""" + global call_count + call_count += 1 + if call_count < REQUIRED_RETRIES: + raise ConnectionError( + f"Movie database unavailable! Attempt {call_count}/{REQUIRED_RETRIES}. Please retry." + ) + return f"{title}: an acclaimed action film, widely regarded as a genre classic." + + +# --------------------------------------------------------------------------- +# Validated structured output +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Rating: + """ + A movie rating, with a score (an integer from 1 to 5) and an explanation. + The explanation MUST mention the score, otherwise it will be rejected as invalid. + """ + + score: typing.Literal[1, 2, 3, 4, 5] + explanation: str + + def __post_init__(self): + if self.score < 1 or self.score > 5: + raise ValueError(f"score must be 1-5, got {self.score}") + if str(self.score) not in self.explanation: + raise ValueError( + f"explanation must mention the score {self.score}, got '{self.explanation}'" + ) + + +# --------------------------------------------------------------------------- +# Template: uses the flaky tool, returns validated structured output +# --------------------------------------------------------------------------- + + +@Template.define +def rate_movie(movie_name: str) -> Rating: + """Look up the movie {movie_name}, then give it a rating.""" + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--movie", type=str, default="Die Hard", help="Movie to rate") + args = parser.parse_args() + + rating = rate_movie(args.movie) + print(f"Rated {args.movie!r} after {call_count} tool attempts:") + print(f"Score: {rating.score}/5") + print(f"Explanation: {rating.explanation}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/flight_booking.py b/docs/source/llm_examples/basics/flight_booking.py new file mode 100644 index 000000000..e3f19dfff --- /dev/null +++ b/docs/source/llm_examples/basics/flight_booking.py @@ -0,0 +1,252 @@ +"""Flight booking with multi-agent delegation. + +Demonstrates: +- Multi-agent delegation: a tool that internally calls a separate + ``@Template.define`` (agent-to-agent delegation) +- Programmatic validation of LLM output with retry +- Interactive human-in-the-loop flow +- ``Agent`` history for conversational seat selection +""" + +import argparse +import dataclasses +import datetime +import enum +from typing import Literal + +from effectful.handlers.llm import Agent, Template + +# --------------------------------------------------------------------------- +# Structured output types +# --------------------------------------------------------------------------- + + +class Airport(enum.StrEnum): + SFO = "SFO" + ANC = "ANC" + FAI = "FAI" + JNU = "JNU" + NYC = "NYC" + LAX = "LAX" + ORD = "ORD" + MIA = "MIA" + BOS = "BOS" + SEA = "SEA" + DFW = "DFW" + DEN = "DEN" + ATL = "ATL" + IAH = "IAH" + + +@dataclasses.dataclass(frozen=True) +class FlightDetails: + flight_number: str + price: int + origin: Airport # three-letter airport code + destination: Airport # three-letter airport code + date: datetime.date # YYYY-MM-DD + + +@dataclasses.dataclass(frozen=True) +class SeatPreference: + """ + User's seat preference extracted from natural language. + + Seats A and F are window seats. Seats C and D are aisle seats. + Row 1 is the front row with extra legroom. + Rows 14 and 20 also have extra legroom. + """ + + row: int # 1-30 + seat: Literal["A", "B", "C", "D", "E", "F"] + + +# --------------------------------------------------------------------------- +# Sample data (in reality, downloaded from a booking site) +# --------------------------------------------------------------------------- + +FLIGHTS_PAGE = """\ +1. Flight SFO-AK123 - $350 - San Francisco (SFO) to Anchorage (ANC) - 2025-01-10 +2. Flight SFO-AK456 - $370 - San Francisco (SFO) to Fairbanks (FAI) - 2025-01-10 +3. Flight SFO-AK789 - $400 - San Francisco (SFO) to Juneau (JNU) - 2025-01-20 +4. Flight NYC-LA101 - $250 - New York (NYC) to Los Angeles (LAX) - 2025-01-10 +5. Flight ORD-MIA202 - $200 - Chicago (ORD) to Miami (MIA) - 2025-01-12 +6. Flight BOS-SEA303 - $120 - Boston (BOS) to Seattle (SEA) - 2025-01-12 +7. Flight DFW-DEN404 - $150 - Dallas (DFW) to Denver (DEN) - 2025-01-10 +8. Flight ATL-IAH505 - $180 - Atlanta (ATL) to Houston (IAH) - 2025-01-10 +""" + +# --------------------------------------------------------------------------- +# Extraction template (inner "agent") +# --------------------------------------------------------------------------- + + +@Template.define +def extract_flights(web_page_text: str) -> list[FlightDetails]: + """Extract all flight details from the following text. + + {web_page_text} + """ + + +# --------------------------------------------------------------------------- +# Flight search agent +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class FlightFinder(Agent): + """Agent that finds flights matching user criteria.""" + + available_flights: list[FlightDetails] + + @Template.define + def find_flight( + self, origin: Airport, destination: Airport, date: datetime.date + ) -> FlightDetails: + """ + Find the cheapest flight from {origin} to {destination} on {date}. + + List of available flights (from the web page): + {self.available_flights} + """ + + +# --------------------------------------------------------------------------- +# Seat selection agent +# --------------------------------------------------------------------------- + + +class SeatSelector(Agent): + """Agent that extracts seat preferences from natural language.""" + + @Template.define + def select_seat(self, user_input: str) -> SeatPreference: + """Extract the user's seat preference from their message. + + {user_input} + """ + + +# --------------------------------------------------------------------------- +# Validation (plain Python, no LLM needed) +# --------------------------------------------------------------------------- + + +def validate_flight( + flight: FlightDetails, origin: Airport, destination: Airport, date: datetime.date +) -> list[str]: + """Check that the selected flight matches the requested criteria.""" + errors = [] + if flight.origin != origin: + errors.append(f"origin should be {origin}, got {flight.origin}") + if flight.destination != destination: + errors.append(f"destination should be {destination}, got {flight.destination}") + if flight.date != date: + errors.append(f"date should be {date}, got {flight.date}") + return errors + + +# --------------------------------------------------------------------------- +# Booking flow +# --------------------------------------------------------------------------- + + +def book_flight( + origin: Airport, + destination: Airport, + date: datetime.date, + interactive: bool = False, + max_retries: int = 3, +) -> None: + """End-to-end flight booking with search, validation, and seat selection.""" + searcher = FlightFinder(available_flights=extract_flights(FLIGHTS_PAGE)) + + # --- Search with validation retry --- + flight = None + for attempt in range(max_retries): + candidate = searcher.find_flight(origin, destination, date) + errors = validate_flight(candidate, origin, destination, date) + if errors: + print(f" [attempt {attempt}] Rejected: {'; '.join(errors)}") + continue + flight = candidate + break + + if flight is None: + print("Could not find a valid flight.") + return + + print( + f" Found: {flight.flight_number} ${flight.price} " + f"({flight.origin}->{flight.destination} on {flight.date})" + ) + + # --- User approval (interactive only) --- + if interactive: + if input(" Book this flight? (yes/no): ").strip().lower() != "yes": + print(" Cancelled.") + return + + # --- Seat selection --- + selector = SeatSelector() + seat_requests = ( + [input(" Seat preference: ")] + if interactive + else ["I'd like a window seat with extra legroom please"] + ) + for request in seat_requests: + seat = selector.select_seat(request) + print(f" Seat: row {seat.row}, seat {seat.seat}") + + print(f" Booked {flight.flight_number}, seat {seat.row}{seat.seat}!") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + airports = list(Airport) + parser.add_argument( + "--origin", + type=Airport, + choices=airports, + default=Airport.SFO, + metavar="CODE", + help="Origin airport code", + ) + parser.add_argument( + "--destination", + type=Airport, + choices=airports, + default=Airport.ANC, + metavar="CODE", + help="Destination airport code", + ) + parser.add_argument( + "--date", + type=datetime.date.fromisoformat, + default=datetime.date(2025, 1, 10), + metavar="YYYY-MM-DD", + help="Travel date (YYYY-MM-DD)", + ) + parser.add_argument( + "--interactive", + action="store_true", + help="Run in interactive mode with user prompts", + ) + args = parser.parse_args() + + book_flight( + origin=args.origin, + destination=args.destination, + date=args.date, + interactive=args.interactive, + ) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/guardrails.py b/docs/source/llm_examples/basics/guardrails.py new file mode 100644 index 000000000..e23f0701d --- /dev/null +++ b/docs/source/llm_examples/basics/guardrails.py @@ -0,0 +1,68 @@ +"""Travel advisor with input guardrails. + +Demonstrates: +- Using one template to validate/guard input before passing it to another +- Simple control-flow gating based on LLM classification +""" + +import argparse + +from effectful.handlers.llm import Template + +# --------------------------------------------------------------------------- +# Templates +# --------------------------------------------------------------------------- + + +@Template.define +def travel_query(user_query: str) -> str: + """ + Produce a concise (<100 word) answer to: {user_query} + """ + + +# --------------------------------------------------------------------------- +# Guarded agent +# --------------------------------------------------------------------------- + + +def answer_travel_query(user_query: str) -> str: + """Only answer travel-related queries; reject everything else.""" + + @Template.define + def is_safe_query(user_query: str) -> bool: + """ + Determine whether the user's query is purely related to travel advice: {user_query} + """ + + if is_safe_query(user_query): + return travel_query(user_query) + else: + return f"Rejected: '{user_query}' is not related to travel advice." + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--queries", + nargs="+", + default=[ + "What are great places to check out in NYC?", + "Should I buy apple stocks?", + ], + metavar="QUERY", + help="User queries to run through the travel-advice guardrail", + ) + args = parser.parse_args() + + for query in args.queries: + print(answer_travel_query(query)) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/hitl.py b/docs/source/llm_examples/basics/hitl.py new file mode 100644 index 000000000..8c9d64aed --- /dev/null +++ b/docs/source/llm_examples/basics/hitl.py @@ -0,0 +1,157 @@ +"""Human-in-the-loop task planner. + +Demonstrates: +- An ``Agent`` that proposes a plan of action steps +- Human approval/rejection of each step before execution +- Feedback from rejection is fed back to the agent via history +- ``@Tool.define`` for executing approved actions +- Non-interactive mode for testing (auto-approves all steps) +""" + +import argparse +import dataclasses +import enum + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# Structured output +# --------------------------------------------------------------------------- + + +class ActionType(enum.StrEnum): + send_email = "send_email" + create_file = "create_file" + schedule_meeting = "schedule_meeting" + done = "done" + + +@dataclasses.dataclass(frozen=True) +class ProposedAction: + action: ActionType + description: str + details: str + + +# --------------------------------------------------------------------------- +# Planner agent +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Planner(Agent): + """Agent that proposes actions one at a time for human approval.""" + + execution_log: list[str] = dataclasses.field(default_factory=list) + + @Tool.define + def execute_action(self, action: ActionType, details: str) -> str: + """Execute an approved action. Returns a confirmation message.""" + msg = f"[executed] {action}: {details}" + self.execution_log.append(msg) + return msg + + @Template.define + def propose_next(self, task: str, feedback: str) -> ProposedAction: + """You are a task planner helping the user accomplish a goal. + + Task: {task} + + Feedback from the last step: {feedback} + + Review the conversation history for previously completed actions. + Propose the next action to take. If the task is complete, + set action to "done". + + If a previous proposal was rejected, propose something different + that addresses the feedback. + """ + + +# --------------------------------------------------------------------------- +# Human-in-the-loop execution +# --------------------------------------------------------------------------- + + +def run_with_approval( + task: str, interactive: bool = False, max_steps: int = 5 +) -> list[str]: + """Run a task planner with human approval for each step.""" + planner = Planner() + feedback = "No actions taken yet. Start planning." + + for step in range(max_steps): + proposal = planner.propose_next(task, feedback) + + if proposal.action == ActionType.done: + print(f" [step {step + 1}] Done: {proposal.description}") + break + + print( + f" [step {step + 1}] Proposed: {proposal.action} - {proposal.description}" + ) + print(f" Details: {proposal.details}") + + if interactive: + answer = input(" Approve? (yes/no + reason): ").strip() + approved = answer.lower().startswith("y") + else: + answer = "yes" + approved = True + + if approved: + result = planner.execute_action(proposal.action, proposal.details) + print(f" {result}") + feedback = f"Approved and executed: {result}" + else: + print(f" [rejected] {answer}") + feedback = f"Rejected: {answer}" + + return list(planner.execution_log) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--interactive", + action="store_true", + help="Run in interactive mode with human approval prompts", + ) + parser.add_argument( + "--max-steps", + type=int, + default=5, + help="Maximum number of action steps", + ) + parser.add_argument( + "--task", + type=str, + default=( + "Organize a team lunch for next Friday. " + "Send an email to the team, create a shared document for " + "restaurant suggestions, and schedule a meeting to finalize plans." + ), + help="The goal for the planner to accomplish", + ) + args = parser.parse_args() + + task = args.task + + print(f"Task: {task}\n") + log = run_with_approval( + task, + interactive=args.interactive, + max_steps=args.max_steps, + ) + print(f"\nExecution log ({len(log)} actions):") + for entry in log: + print(f" {entry}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/image_input.py b/docs/source/llm_examples/basics/image_input.py new file mode 100644 index 000000000..68f0d5bee --- /dev/null +++ b/docs/source/llm_examples/basics/image_input.py @@ -0,0 +1,50 @@ +"""Passing PIL images directly to a template. + +Demonstrates: +- Templates accepting ``PIL.Image.Image`` arguments +- Inline base64 image data so the script is self-contained +""" + +import argparse +import base64 +import io + +from PIL import Image + +from effectful.handlers.llm import Template + + +@Template.define +def describe_image(image: Image.Image) -> str: + """Return a short description of the following image. + {image} + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--image", + type=str, + default=None, + metavar="PATH", + help="Path to an image file to describe (defaults to a built-in 32x32 smiley face)", + ) + args = parser.parse_args() + + if args.image is not None: + image = Image.open(args.image) + else: + IMAGE_BASE64 = ( + "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAAhElEQVR4nO2W4QqA" + "MAiEVXr/VzYWDGoMdk7Cgrt/sUs/DqZTd3EplFU2JwATYAJMoOlAB4bq89s95+Mg" + "+gyAchsKAYplBBBA43hFhfxnUixDjdEUUL8hpr7R0KLdt9qElzcyiu8As+Kr8zQA" + "mgLavAl+kIzFZyCRxtsAmWb/voZvqRzgBE1sIDuVFX4eAAAAAElFTkSuQmCC" + ) + image = Image.open(io.BytesIO(base64.b64decode(IMAGE_BASE64))) + + print(describe_image(image)) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/image_tool.py b/docs/source/llm_examples/basics/image_tool.py new file mode 100644 index 000000000..abcff550e --- /dev/null +++ b/docs/source/llm_examples/basics/image_tool.py @@ -0,0 +1,97 @@ +import argparse +import pathlib + +from PIL import Image + +from effectful.handlers.llm import Agent, Template, Tool + + +class ImageTools(Agent): + """You are an image processing agent.""" + + _image_to_handle: dict[int, int] + _handle_to_image: dict[int, Image.Image] + + def __init__(self): + self._image_to_handle = {} + self._handle_to_image = {} + + def _encode(self, image: Image.Image) -> int: + image_id = id(image) + handle = self._image_to_handle.get(image_id, None) + if handle is not None: + return handle + + handle = len(self._image_to_handle) + self._image_to_handle[image_id] = handle + + assert handle not in self._handle_to_image + self._handle_to_image[handle] = image + return handle + + def _decode(self, image_handle: int) -> Image.Image: + return self._handle_to_image[image_handle] + + @Tool.define + def rotate(self, image: int, angle: float) -> int: + """Returns a rotated copy of this image. The copy is rotated by `angle` + degrees counterclockwise around the image center. + + """ + return self._encode(self._decode(image).rotate(angle)) + + @Tool.define + def concat_horiz(self, i1_h: int, i2_h: int) -> int: + """Concatenates two images horizontally. The larger image will be + cropped to the height of the smaller image. + + """ + i1 = self._decode(i1_h) + i2 = self._decode(i2_h) + i3 = Image.new("RGB", (i1.width + i2.width, min(i1.height, i2.height))) + i3.paste(i1, (0, 0)) + i3.paste(i2, (i1.width, 0)) + return self._encode(i3) + + @Template.define + def _rotate_and_concat(self, i: int) -> int: + """Create an image consisting of four copies of the image {i} + concatenated horizontally. Each copy should be rotated 90 degrees from + the previous. + + """ + + def rotate_and_concat(self, i: Image.Image) -> Image.Image: + return self._decode(self._rotate_and_concat(self._encode(i))) + + +def main() -> None: + # The shared static directory is ``docs/source/_static``; this file lives at + # ``docs/source/llm_examples/basics/``, so it is three levels up -- it was two + # before the examples moved into ``basics/``, which left this default pointing + # at a path that does not exist. + DEFAULT_IMAGE = ( + pathlib.Path(__file__).resolve().parents[2] + / "_static" + / "img" + / "chirho_logo_wide.png" + ) + + parser = argparse.ArgumentParser(description=__doc__ or ImageTools.__doc__) + parser.add_argument( + "--image", + type=str, + default=str(DEFAULT_IMAGE), + metavar="PATH", + help="Path to the input image to rotate-and-concatenate.", + ) + args = parser.parse_args() + + image_agent = ImageTools() + img = Image.open(args.image) + + image_agent.rotate_and_concat(img).show() + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/lexical_scope.py b/docs/source/llm_examples/basics/lexical_scope.py new file mode 100644 index 000000000..697205218 --- /dev/null +++ b/docs/source/llm_examples/basics/lexical_scope.py @@ -0,0 +1,99 @@ +"""Composition via lexical scope: auto-captured sub-templates, invoked two ways. + +Demonstrates: +- Module-level @Template.define sub-templates auto-captured into other templates' + lexical scope, with no explicit registration +- An Agent grouping @Tool.define tools with a @Template.define orchestrator that + calls those tools and the sub-templates directly (model-driven composition) +- A template returning a Callable: the model synthesizes a function that calls the + same sub-templates when run (code-driven composition), via the eval provider +- inspect.getsource on the synthesized function +""" + +import argparse +import inspect +from collections.abc import Callable +from typing import Literal + +from effectful.handlers.llm import Agent, Template, Tool + + +@Template.define +def story_with_moral(topic: str) -> str: + """Write a short story about {topic} and end with a moral lesson. Do not use any tools.""" + + +@Template.define +def story_funny(topic: str) -> str: + """Write a funny, humorous story about {topic}. Do not use any tools.""" + + +class TripPlanner(Agent): + """Plans a trip to a city with good weather and tells a story about visiting it.""" + + @Tool.define + def cities(self) -> list[str]: + """Return a list of candidate destination cities.""" + return ["Chicago", "New York", "Barcelona"] + + @Tool.define + def weather(self, city: str) -> str: + """Given a city name, return a short description of its weather.""" + status = {"Chicago": "cold", "New York": "wet", "Barcelona": "sunny"} + return status.get(city, "unknown") + + @Template.define + def plan_trip_story(self, style: str) -> str: + """Use the relevant tools to identify a city that has good (sunny) + weather. Then write a short story about visiting that city in the requested + style: {style}""" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--style", + type=str, + choices=["moral", "funny"], + default="funny", + help="Style of the story to produce", + ) + parser.add_argument( + "--topic", + type=str, + default="a curious cat", + help="Topic for the synthesized story function to run on", + ) + parser.add_argument( + "--method", + type=str, + choices=["model", "code"], + default="model", + help="Whether to run the model-driven or code-driven composition", + ) + args = parser.parse_args() + + if args.method == "model": + # (1) Model-driven: the orchestrator template calls tools and sub-templates. + print("=== Orchestrator template (model-driven composition) ===") + planner = TripPlanner() + print(planner.plan_trip_story(args.style)) + + elif args.method == "code": + + @Template.define + def write_story_fn(style: Literal["moral", "funny"]) -> Callable[[str], str]: + """Generate a Python function that takes a topic string and returns a story + about it in the {style} style. The function should delegate the writing to the + `story_funny` sub-template for humor, or `story_with_moral` for a lesson.""" + + # (2) Code-driven: the model synthesizes a function that calls the sub-templates. + print(f"\n=== Synthesized higher-order function (style={args.style}) ===") + story_fn = write_story_fn(args.style) + print(inspect.getsource(story_fn)) + print(f"\n=== Running it on {args.topic!r} ===") + print(story_fn(args.topic)) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/map_reduce.py b/docs/source/llm_examples/basics/map_reduce.py new file mode 100644 index 000000000..3efca2697 --- /dev/null +++ b/docs/source/llm_examples/basics/map_reduce.py @@ -0,0 +1,138 @@ +"""Map-reduce resume evaluation. + +Demonstrates: +- Fan-out: evaluating multiple items independently with the same template +- Reduce: aggregating individual results into a summary +- ``asyncio.gather`` with ``asyncio.to_thread`` for parallel LLM calls +- Structured output with dataclasses +""" + +import argparse +import asyncio +import collections.abc +import dataclasses +import functools +import typing + +from effectful.handlers.llm import Template + +# --------------------------------------------------------------------------- +# Structured output +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class Evaluation: + name: str + qualified: bool + strengths: str + weaknesses: str + score: typing.Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + +# --------------------------------------------------------------------------- +# Templates +# --------------------------------------------------------------------------- + + +@Template.define +def evaluate_resume(resume: str, job_description: str) -> Evaluation: + """You are a hiring manager. Evaluate this resume against the job + description and produce a structured evaluation. + + Job description: {job_description} + + Resume: + {resume} + """ + + +@Template.define +def summarize_evaluations( + job_description: str, + evaluations: collections.abc.Sequence[Evaluation], +) -> str: + """You are a hiring manager summarizing candidate evaluations. + + Job description: {job_description} + + Individual evaluations: + {evaluations} + + Provide a brief summary: rank the candidates from best to worst, + highlight the top candidate, and note any concerns. + """ + + +# --------------------------------------------------------------------------- +# Sample data +# --------------------------------------------------------------------------- + +JOB_DESCRIPTION = ( + "Senior Python Developer: 5+ years Python experience, " + "familiarity with web frameworks (Django/Flask), " + "database design, and cloud deployment (AWS/GCP)." +) + +RESUMES = [ + "Alice Chen - 7 years Python, Django expert, AWS certified, " + "led team of 5, built microservices architecture at FinTech startup.", + "Bob Smith - 3 years Python, 2 years JavaScript, some Flask experience, " + "junior developer at small agency, strong communication skills.", + "Carol Davis - 10 years software engineering, 6 years Python, " + "GCP specialist, PostgreSQL expert, open-source contributor, " + "previously senior engineer at Google.", + "Dave Wilson - 4 years Python, self-taught, built several side projects, " + "no professional experience with web frameworks or cloud platforms.", +] + +# --------------------------------------------------------------------------- +# Map-reduce pipeline +# --------------------------------------------------------------------------- + + +async def map_reduce_evaluate( + resumes: list[str], + job_description: str, +) -> str: + """Evaluate resumes in parallel (map), then summarize (reduce).""" + # Map: fork/join -- evaluate each resume concurrently via asyncio.gather + + # asyncio.to_thread (sync template calls run in parallel threads). + evaluate = functools.partial(asyncio.to_thread, evaluate_resume) + evaluations: list[Evaluation] = list( + await asyncio.gather(*(evaluate(resume, job_description) for resume in resumes)) + ) + + # Print individual evaluations + for ev in evaluations: + print(f" {ev.name}: score={ev.score}/10, qualified={ev.qualified}") + print(f" + {ev.strengths}") + print(f" - {ev.weaknesses}") + + # Reduce: summarize all evaluations + return summarize_evaluations(job_description, evaluations) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--job-description", + type=str, + metavar="TEXT", + default=JOB_DESCRIPTION, + help="Job description to evaluate resumes against.", + ) + args = parser.parse_args() + + print(f"Evaluating {len(RESUMES)} resumes for: {args.job_description}\n") + summary = asyncio.run(map_reduce_evaluate(RESUMES, args.job_description)) + print(f"\n{summary}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/rag.py b/docs/source/llm_examples/basics/rag.py new file mode 100644 index 000000000..7e4cab7f9 --- /dev/null +++ b/docs/source/llm_examples/basics/rag.py @@ -0,0 +1,185 @@ +"""Retrieval-augmented generation (RAG). + +Demonstrates: +- Offline: chunking documents, embedding, and indexing +- Online: embedding a query, retrieving relevant chunks, and generating + a grounded answer +- ``@Tool.define`` to expose retrieval as a tool the LLM can call +- Separation of indexing (plain Python) from generation (``@Template.define``) +""" + +import argparse +import dataclasses + +import litellm +import numpy as np + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# Vector index +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class VectorIndex: + """Simple in-memory vector index using L2 distance.""" + + model: str + chunks: list[str] = dataclasses.field(default_factory=list) + embeddings: list[np.ndarray] = dataclasses.field(default_factory=list) + + def get_embedding(self, text: str) -> np.ndarray: + """Get an embedding vector for the given text using litellm.""" + response = litellm.embedding(model=self.model, input=text) + return np.array(response.data[0]["embedding"], dtype=np.float32) + + def add(self, text: str) -> None: + """Add a text chunk to the index.""" + self.chunks.append(text) + self.embeddings.append(self.get_embedding(text)) + + def search(self, query: str, top_k: int = 3) -> list[str]: + """Return the top-k most similar chunks to the query.""" + if not self.embeddings: + return [] + query_emb = self.get_embedding(query) + distances = [float(((emb - query_emb) ** 2).sum()) for emb in self.embeddings] + indices = sorted(range(len(distances)), key=lambda i: distances[i]) + return [self.chunks[i] for i in indices[:top_k]] + + +# --------------------------------------------------------------------------- +# Chunking +# --------------------------------------------------------------------------- + + +def chunk_text(text: str, chunk_size: int = 200, overlap: int = 50) -> list[str]: + """Split text into overlapping word-level chunks.""" + words = text.split() + chunks = [] + start = 0 + while start < len(words): + end = start + chunk_size + chunks.append(" ".join(words[start:end])) + start += chunk_size - overlap + return chunks + + +# --------------------------------------------------------------------------- +# Sample documents +# --------------------------------------------------------------------------- + +DOCUMENTS = [ + """The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars + in Paris, France. It is named after the engineer Gustave Eiffel, whose + company designed and built the tower from 1887 to 1889 as the centerpiece + of the 1889 World's Fair. Although initially criticized by some of France's + leading artists and intellectuals, the tower has become a global icon of + France and one of the most recognizable structures in the world. The tower + is 330 metres tall, about the same height as an 81-storey building, and + is the tallest structure in Paris. It was the first structure in the world + to reach a height of 300 metres.""", + """The Great Wall of China is a series of fortifications that were built + across the historical northern borders of ancient Chinese states and + Imperial China as protection against various nomadic groups. The total + length of all sections ever built is more than 20,000 km. Several walls + were built from as early as the 7th century BC, with selective stretches + later joined together by Qin Shi Huang, the first emperor of China. The + best-preserved sections of the wall date from the Ming dynasty + (1368-1644). The wall's purpose was defensive, and it featured + watchtowers, troop barracks, and signaling capabilities.""", + """The Colosseum, also known as the Flavian Amphitheatre, is an oval + amphitheatre in the centre of the city of Rome, Italy. It is the largest + ancient amphitheatre ever built, and is still the largest standing + amphitheatre in the world, despite its age. Construction began under + the emperor Vespasian in AD 72 and was completed in AD 80 under his + successor and heir, Titus. The Colosseum could hold an estimated 50,000 + to 80,000 spectators at various points in its history, and was used for + gladiatorial contests and public spectacles including animal hunts, + executions, re-enactments of famous battles, and dramas.""", +] + +# --------------------------------------------------------------------------- +# Build the index (offline phase) +# --------------------------------------------------------------------------- + + +def build_index(documents: list[str], embedding_model: str) -> VectorIndex: + """Chunk and index a collection of documents.""" + index = VectorIndex(model=embedding_model) + for doc in documents: + for chunk in chunk_text(doc, chunk_size=60, overlap=15): + index.add(chunk) + print(f"Indexed {len(index.chunks)} chunks from {len(documents)} documents") + return index + + +# --------------------------------------------------------------------------- +# RAG agent (online phase): the `retrieve` tool and `answer_question` template +# share one instance, so the tool is auto-captured from lexical scope. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class RAGAgent(Agent): + """Answers a question grounded in the vector index via a retrieval tool.""" + + index: VectorIndex + + @Tool.define + def retrieve(self, query: str, top_k: int = 3) -> list[str]: + """Return the top-k most similar chunks to the query.""" + return self.index.search(query, top_k) + + @Template.define + def answer_question(self, question: str) -> str: + """You are a helpful assistant. Answer the user's question using ONLY + information retrieved from the knowledge base via the retrieve tool. + + If the retrieved information doesn't contain the answer, say so. + Always cite which document your information comes from. + + Question: {question} + """ + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--embedding-model", + type=str, + default="lm_studio/text-embedding-embeddinggemma-300m-qat", + help="Embedding model to use", + ) + parser.add_argument( + "--questions", + type=str, + nargs="+", + metavar="QUESTION", + default=[ + "How tall is the Eiffel Tower?", + "When was the Great Wall of China built?", + "How many spectators could the Colosseum hold?", + ], + help="Questions to answer against the indexed documents", + ) + args = parser.parse_args() + + # Offline: build the index once. + index = build_index(DOCUMENTS, embedding_model=args.embedding_model) + + # Online: answer each question with a fresh (stateless) agent over that index. + for question in args.questions: + print(f"\nQ: {question}") + answer = RAGAgent(index=index).answer_question(question) + print(f"A: {answer}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/research_agent.py b/docs/source/llm_examples/basics/research_agent.py new file mode 100644 index 000000000..76d69a736 --- /dev/null +++ b/docs/source/llm_examples/basics/research_agent.py @@ -0,0 +1,162 @@ +"""Research agent with web search and LLM quality control. + +Demonstrates: +- @Tool.define web-search tool, auto-captured into templates from lexical scope +- An Agent subclass with persistent conversation history +- One Template judging another's output, returning a structured QualityJudgment + (a bool plus written feedback) +- A feedback-driven refinement loop: answer -> judge -> refine -> judge -> ... +""" + +import argparse +import dataclasses +import urllib.parse + +import requests + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# Search tool +# --------------------------------------------------------------------------- + + +@Tool.define +def search_web(query: str) -> str: + """Search Wikipedia for a topic and return a summary. The query can be a topic name or a natural language question.""" + search_url = "https://en.wikipedia.org/w/api.php?" + urllib.parse.urlencode( + { + "action": "query", + "list": "search", + "srsearch": query, + "srlimit": 1, + "format": "json", + } + ) + search_data = requests.get( + search_url, headers={"User-Agent": "effectful-example/1.0"} + ).json() + results = search_data.get("query", {}).get("search", []) + if not results: + return f"No results found for: {query}" + title = results[0]["title"] + + summary_url = "https://en.wikipedia.org/w/api.php?" + urllib.parse.urlencode( + { + "action": "query", + "titles": title, + "prop": "extracts", + "exintro": True, + "explaintext": True, + "format": "json", + } + ) + summary_data = requests.get( + summary_url, headers={"User-Agent": "effectful-example/1.0"} + ).json() + page = next(iter(summary_data["query"]["pages"].values())) + extract = page.get("extract", "No summary available.") + url = f"https://en.wikipedia.org/wiki/{urllib.parse.quote(title.replace(' ', '_'))}" + + return f"# {title}\n\n{extract}\n\nSource: {url}" + + +# --------------------------------------------------------------------------- +# Structured output for quality judgment +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class QualityJudgment: + is_acceptable: bool + feedback: str + + +# --------------------------------------------------------------------------- +# Research agent (persistent history; search_web auto-captured from scope) +# --------------------------------------------------------------------------- + + +class Researcher(Agent): + """Agent that answers research questions using web search, refining on feedback.""" + + @Template.define + def answer(self, question: str) -> str: + """You are a research assistant. Use the search tool to find accurate, + specific information, then answer the question: {question}""" + + @Template.define + def refine(self, question: str, feedback: str) -> str: + """A reviewer rejected your previous answer to the question ({question}) + with this feedback: {feedback}. Use the search tool as needed and provide + an improved answer that addresses the feedback.""" + + +# --------------------------------------------------------------------------- +# Supervisor (quality judge) +# --------------------------------------------------------------------------- + + +@Template.define +def judge_quality(question: str, answer: str) -> QualityJudgment: + """You are a strict quality reviewer. Evaluate whether this answer adequately + addresses the question with accurate, specific information. + + Question: {question} + Answer: {answer} + + An answer is acceptable if it contains specific facts (names, dates, numbers) + relevant to the question. Vague or generic answers should be rejected; when + rejecting, explain in the feedback what is missing. + """ + + +# --------------------------------------------------------------------------- +# Supervised agent loop +# --------------------------------------------------------------------------- + + +def research_agent(question: str, max_retries: int = 3) -> str: + """Answer a question, refining on supervisor feedback until it is acceptable.""" + researcher = Researcher() + answer = researcher.answer(question) + + for attempt in range(1, max_retries + 1): + judgment = judge_quality(question, answer) + if judgment.is_acceptable: + print(f"[supervisor] Accepted on attempt {attempt}") + return answer + print(f"[supervisor] Rejected attempt {attempt}: {judgment.feedback}") + answer = researcher.refine(question, judgment.feedback) + + print("[supervisor] Returning best effort after max retries") + return answer + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--question", + type=str, + default="What year was the Eiffel Tower completed and how tall is it?", + help="The question to research", + ) + parser.add_argument( + "--max-retries", + type=int, + default=3, + help="Maximum number of supervisor rejections before returning best effort", + ) + args = parser.parse_args() + + result = research_agent(args.question, max_retries=args.max_retries) + print(f"\nFinal answer: {result}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/basics/text2sql.py b/docs/source/llm_examples/basics/text2sql.py new file mode 100644 index 000000000..528483982 --- /dev/null +++ b/docs/source/llm_examples/basics/text2sql.py @@ -0,0 +1,156 @@ +"""Natural language to SQL with LLM-powered debug loop. + +Demonstrates: +- Generating SQL from natural language using ``@Template.define`` +- Executing SQL against a real SQLite database +- Feeding execution errors back to the LLM for iterative fixing +- ``@Tool.define`` to expose the database schema as a tool +""" + +import argparse +import sqlite3 +import textwrap + +from effectful.handlers.llm import Template + +# --------------------------------------------------------------------------- +# In-memory database setup +# --------------------------------------------------------------------------- + + +def create_sample_db() -> sqlite3.Connection: + """Create a sample SQLite database with employee data.""" + conn = sqlite3.connect(":memory:") + conn.executescript( + textwrap.dedent("""\ + CREATE TABLE departments ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + budget REAL NOT NULL + ); + CREATE TABLE employees ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + department_id INTEGER REFERENCES departments(id), + salary REAL NOT NULL, + hire_date TEXT NOT NULL + ); + INSERT INTO departments VALUES (1, 'Engineering', 500000); + INSERT INTO departments VALUES (2, 'Marketing', 200000); + INSERT INTO departments VALUES (3, 'Sales', 300000); + INSERT INTO employees VALUES (1, 'Alice', 1, 120000, '2020-01-15'); + INSERT INTO employees VALUES (2, 'Bob', 1, 110000, '2021-03-22'); + INSERT INTO employees VALUES (3, 'Carol', 2, 95000, '2019-07-01'); + INSERT INTO employees VALUES (4, 'Dave', 3, 105000, '2022-11-10'); + INSERT INTO employees VALUES (5, 'Eve', 1, 130000, '2018-05-20'); + INSERT INTO employees VALUES (6, 'Frank', 3, 98000, '2023-01-05'); + """) + ) + return conn + + +def get_schema(conn: sqlite3.Connection) -> str: + """Extract the schema from a SQLite database.""" + cursor = conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' ORDER BY name" + ) + return "\n\n".join(row[0] for row in cursor if row[0]) + + +# --------------------------------------------------------------------------- +# Templates +# --------------------------------------------------------------------------- + + +@Template.define +def generate_sql(question: str, db_schema: str) -> str: + """You are a SQL expert. Given this database schema: + + {db_schema} + + Write a SQLite query that answers: {question} + + Return ONLY the SQL query, no explanation. + """ + + +@Template.define +def fix_sql(question: str, db_schema: str, bad_sql: str, error: str) -> str: + """You are a SQL expert. Your previous query had an error. + + Database schema: + {db_schema} + + Original question: {question} + Failed SQL: {bad_sql} + Error: {error} + + Write a corrected SQLite query. Return ONLY the SQL query. + """ + + +# --------------------------------------------------------------------------- +# Text-to-SQL agent with debug loop +# --------------------------------------------------------------------------- + + +def text_to_sql( + conn: sqlite3.Connection, question: str, max_retries: int = 3 +) -> list[tuple]: + """Convert a natural language question to SQL and execute it. + + If the query fails, feed the error back to the LLM to fix it, + up to ``max_retries`` times. + """ + schema = get_schema(conn) + sql = generate_sql(question, schema) + + for attempt in range(max_retries + 1): + # Strip markdown fences if the LLM wraps the SQL + clean_sql = sql.strip().removeprefix("```sql").removesuffix("```").strip() + print(f" [attempt {attempt + 1}] {clean_sql}") + + try: + cursor = conn.execute(clean_sql) + return cursor.fetchall() + except Exception as e: + if attempt < max_retries: + print(f" [error] {e}") + sql = fix_sql(question, schema, clean_sql, str(e)) + else: + raise + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--questions", + nargs="+", + metavar="QUESTION", + default=[ + "What is the average salary by department?", + "Who is the highest paid employee?", + "How many employees were hired after 2021?", + ], + help="Natural-language questions to answer against the sample database", + ) + args = parser.parse_args() + + conn = create_sample_db() + for question in args.questions: + print(f"\nQ: {question}") + try: + rows = text_to_sql(conn, question) + for row in rows: + print(f" => {row}") + except Exception as e: + print(f" FAILED: {e}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/choreographies/__init__.py b/docs/source/llm_examples/choreographies/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/choreographies/library.py b/docs/source/llm_examples/choreographies/library.py new file mode 100644 index 000000000..c34bece11 --- /dev/null +++ b/docs/source/llm_examples/choreographies/library.py @@ -0,0 +1,608 @@ +"""Choreographic programming for multi-agent LLM systems. + +Write a single ``async`` function describing how agents interact from a global +perspective, then run it with automatic endpoint projection (EPP). Every agent +runs that same function as its own `asyncio.Task`, and inter-agent +communication falls out of ordinary asyncio primitives. + +## How it works + +Each `step` in the choreography is assigned an incrementing step ID. Because +every agent runs the same program, and every step's result is shared, all +agents allocate the same IDs in the same order. Each ID names an +`asyncio.Future`; for a given step, `EndpointProjection` either: + +- **executes** it and resolves the future, if the step's template belongs to + this agent; or +- **awaits** that future, if it belongs to another agent. + +That is the whole coordination mechanism. A future is exactly a write-once, +read-by-many cell, which is what a step result is: the architect computes step +0 once, and every other agent reads it. Waiting is event-driven -- nothing +polls, and there is no interval to tune. + +`scatter` needs the one thing futures don't provide, namely handing each item +to exactly one of several workers. That is a work queue, so it uses one: +an `asyncio.Queue` of item indices that the agents in the pool drain with +`get_nowait`. Whoever is free takes the next item, which balances load by +construction. + +Results live in memory, so by default an interrupted run starts over. Give +`Choreography` a *log* path and each step is written to SQLite as it completes; +a later run over the same path replays those results and resumes at the first +step that never finished. (Agent *history* is a separate matter -- give an +`~effectful.handlers.llm.template.Agent` an ``agent_id`` and install +`~effectful.handlers.llm.completions.SQLitePersister` to checkpoint it.) + +## Why the program is async, and where the threads went + +`effectful`'s handler stack is synchronous, top to bottom: `Template.__apply__`, +`fwd`, and everything in `effectful.handlers.llm.completions` down to +`litellm.completion` are blocking calls. Two consequences shape this module. + +*A handler's body must run synchronously -- but it may return an awaitable.* +`coproduct` wraps every handler in a synchronous continuation (see +`effectful.internals.runtime._set_prompt`), and `Operation.__call__` binds +`~effectful.ops.semantics.fwd` around the call itself, so an ``async def`` +handler would return an un-awaited coroutine whose body later ran outside both +bindings. `step` and `scatter` are therefore ordinary `Operation`s whose +implementations return coroutines rather than being coroutines, which is also +what lets a step ID be allocated while `step` is being called. It is why the +choreography spells its steps out with ``await step(...)`` rather than calling +``architect.plan(spec)`` directly. + +*A template call must still run on a thread.* `step` hands the blocking call to +a worker thread and awaits it, so an agent waiting on a peer costs a suspended +coroutine rather than a parked thread. Note that the naive alternative -- +wrapping each agent's whole program in `asyncio.to_thread` -- deadlocks here: +agents block on each other, `asyncio.to_thread` draws from a default executor +of ``min(32, cpu_count + 4)`` workers, and waiting agents hold workers that the +agents they wait for can never get. `Choreography` sizes its own executor to +the number of agents for the same reason. + +## Primitives + +`step` + One template call: executed by its owner, shared with everyone else. Inside + a `scatter` item, where the item is already the step, it is just the call. +`scatter` + Distribute items across a pool of same-role agents, each item going to + whichever agent is free. + +Several scatters run concurrently with `asyncio.gather`; agents belonging to +more than one group work on all of them at once:: + + specs, tests, proofs = await asyncio.gather( + scatter(blocks, spec_writer, lambda w, b: step(w.write_spec, b)), + scatter(blocks, tester, lambda t, b: step(t.write_tests, b)), + scatter(blocks, prover, lambda p, b: step(p.prove, b)), + ) + +Step IDs are allocated when `step`/`scatter` is *called*, not when the returned +awaitable is *awaited*, so the IDs in a `asyncio.gather` are deterministic and +agree across agents. + +## Writing one + +A choreography is an ``async`` function whose parameters are the agents. It +reads as the workflow, from nobody's point of view in particular:: + + async def build_codebase(project_spec, architect, coder, reviewer): + plan = await step(architect.plan_modules, project_spec) + codes = await scatter( + plan["modules"], coder, + lambda c, mod: step(c.implement_module, str(mod)), + ) + return [await step(reviewer.review_code, code) for code in codes] + +Hand it the agents and run it. A role may be filled by several agents, which +is what gives `scatter` a pool to hand work to:: + + choreo = Choreography( + build_codebase, + agents=[architect, coder1, coder2, reviewer], + log="./state/steps.db", # optional; resume where an earlier run stopped + ) + with handler(LiteLLMProvider(model="gpt-4o-mini")), handler(RetryLLMHandler()): + reviews = choreo( + "Build a URL slugify library", + architect=architect, + coder=[coder1, coder2], + reviewer=reviewer, + ) + +``multi_agent_choreography.py``, alongside this module, is a complete, runnable +version: agents with tools, a review-and-fix loop, and resumption. + +""" + +import asyncio +import concurrent.futures +import contextlib +import contextvars +import functools +import os +import pathlib +import pickle +import sqlite3 +import typing +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +from effectful.handlers.llm.types import Agent +from effectful.ops.semantics import handler +from effectful.ops.syntax import ObjectInterpretation, implements +from effectful.ops.types import Operation + + +class ChoreographyError(Exception): + """Raised when a choreography fails because one of its agents failed.""" + + +# ── Shared step state ───────────────────────────────────────────── + + +class _Steps: + """The shared state of one choreography run. + + Two dictionaries, keyed by step ID: a `asyncio.Future` per step, holding + the result its owner computes and every other agent awaits, and a + `asyncio.Queue` per scatter, holding the item indices its pool drains. + + Both accessors are get-or-create, and neither awaits, so concurrent agents + cannot interleave inside them: whichever agent reaches a step first creates + its cell and the rest find it. That also means one of these belongs to one + event loop, which is why `Choreography` makes a fresh one per run. + + Given the path to a log, each step is also written to SQLite as it + resolves, and `replay` reads them back at the start of a later run. The + file is the whole of that durable state -- these objects come and go with + the runs that use them. + """ + + def __init__(self, log: pathlib.Path | None = None) -> None: + self._results: dict[str, asyncio.Future] = {} + self._work: dict[str, asyncio.Queue[int]] = {} + self._log = log + if log is not None: + # Once per run, rather than on every step that gets recorded. + log.parent.mkdir(parents=True, exist_ok=True) + with self._connect() as conn: + conn.execute( + "CREATE TABLE IF NOT EXISTS steps " + "(id TEXT PRIMARY KEY, result BLOB NOT NULL)" + ) + + def _connect(self) -> contextlib.AbstractContextManager[sqlite3.Connection]: + """A connection to the log, closing on exit, in autocommit mode. + + Autocommit because every write is a single statement: there is nothing + to group into a transaction, and a step's result is durable the moment + it is written. + """ + conn = sqlite3.connect(str(self._log), isolation_level=None) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + return contextlib.closing(conn) + + def result(self, step_id: str) -> asyncio.Future: + """The future holding *step_id*'s result.""" + future = self._results.get(step_id) + if future is None: + future = self._results[step_id] = asyncio.get_running_loop().create_future() + return future + + def resolve(self, step_id: str, value: Any) -> None: + """Publish *value* as *step_id*'s result, recording it first. + + Recording before publishing keeps the log ahead of the run: a crash + between the two costs one step's re-execution on the next run, whereas + the other order would report a step as done that no later run knows + about. + """ + if self._log is not None: + with self._connect() as conn: + conn.execute( + "INSERT OR REPLACE INTO steps (id, result) VALUES (?, ?)", + (step_id, pickle.dumps(value)), + ) + self.result(step_id).set_result(value) + + def replay(self) -> int: + """Pre-resolve the steps an earlier run recorded, and return how many.""" + if self._log is None: + return 0 + with self._connect() as conn: + rows = conn.execute("SELECT id, result FROM steps").fetchall() + for step_id, blob in rows: + future = self.result(step_id) + if not future.done(): + future.set_result(pickle.loads(blob)) + return len(rows) + + def work(self, step_id: str, results: Sequence[asyncio.Future]) -> asyncio.Queue: + """The queue of item indices for the scatter at *step_id*. + + Items already resolved -- replayed from a previous run -- are left out, + so a resumed scatter only distributes what is still outstanding. + """ + queue = self._work.get(step_id) + if queue is None: + queue = self._work[step_id] = asyncio.Queue() + for index, result in enumerate(results): + if not result.done(): + queue.put_nowait(index) + return queue + + +def _fail(future: asyncio.Future, error: BaseException) -> None: + """Fail *future*, so agents awaiting it see the error instead of hanging.""" + if future.done(): + return + future.set_exception(error) + # The agents that would have retrieved this are normally cancelled by the + # task group before they get the chance, and asyncio complains at + # collection time about an exception nobody read. Read it here: the error + # still reaches any live waiter, and the failure is reported by the agent + # that actually raised it. + future.exception() + + +# ── Endpoint projection ─────────────────────────────────────────── + + +@Operation.define +def step[**P, T]( + template: Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> Awaitable[T]: + """Take one step of a choreography, and return an awaitable for its result. + + Under `EndpointProjection`, the agent that owns *template* executes the + step while the others await its result; a step recorded by an earlier run + (see `Choreography`'s *log*) returns without calling the model at all. A + template bound + to no agent is executed by every agent. + + The step ID is allocated when `step` is called, not when its result is + awaited, so concurrent steps still get the same IDs in the same order on + every agent. + + Inside `scatter` the enclosing item *is* the step -- it has its own ID and + its own entry in the log -- so a step there is simply the call itself, run + on the choreography's thread pool with no further bookkeeping. A step on + another agent's template is refused, since a scatter item is work one agent + took on alone. + + Unhandled -- outside any choreography -- this is `asyncio.to_thread`, so a + choreographic program is still runnable on its own, one step after another: + + >>> import asyncio + >>> asyncio.run(step(str.upper, "a step is a call, until it is projected")) + 'A STEP IS A CALL, UNTIL IT IS PROJECTED' + """ + return asyncio.to_thread(template, *args, **kwargs) + + +@Operation.define +def scatter[A: Agent, T, U]( + items: Sequence[T], + agent: A | Sequence[A], + fn: Callable[[A, T], Awaitable[U]], +) -> Awaitable[list[U]]: + """Distribute *items* over *agent* by calling ``await fn(agent, item)``. + + *agent* may be a single agent or a pool of same-role agents. Under + `EndpointProjection` the pool draws items from a shared `asyncio.Queue` + until it is empty, which balances load by construction: a fast agent takes + more items. + + Results come back in *items* order, whoever computed them. + + Unhandled, items are processed sequentially, round-robin over the pool. + + *fn* should only touch the agent it is handed; a `step` inside it on any + other agent's template is refused. + """ + return _scatter_sequentially(items, agent, fn) + + +async def _scatter_sequentially[A: Agent, T, U]( + items: Sequence[T], + agent: A | Sequence[A], + fn: Callable[[A, T], Awaitable[U]], +) -> list[U]: + agents = [agent] if isinstance(agent, Agent) else list(agent) + return [await fn(agents[i % len(agents)], item) for i, item in enumerate(items)] + + +class EndpointProjection(ObjectInterpretation): + """Projects a choreographic program onto a single agent. + + `Choreography` installs one per agent, and that is what makes the agents + -- all running the same program -- behave differently: `step` and + `scatter` route through the projection for the task it was installed in. + + Each implementation runs synchronously and *returns* an awaitable rather + than being a coroutine function itself. That is what keeps step IDs in + lockstep, since the ID is allocated while `step` is being called, and it + is also what keeps `~effectful.ops.semantics.fwd` meaningful: `effectful` + binds it around the synchronous call, so a handler that returned an + un-awaited coroutine would run its body after that binding was gone. + + Args: + agent: The agent this projection speaks for. + steps: The run's shared step state. Every agent in a choreography must + be given the same one -- it is how they exchange results. + agent_ids: The IDs of every agent in the run, used to reject a step + belonging to an agent that is not participating. ``None`` skips + the check. + executor: Thread pool for blocking template calls. ``None`` uses + asyncio's default executor, which is only safe when agents do not + wait on each other -- `Choreography` always passes its own. + """ + + def __init__( + self, + agent: Agent, + steps: "_Steps", + agent_ids: frozenset[str] | None = None, + executor: concurrent.futures.Executor | None = None, + ) -> None: + self._agent = agent + self._agent_id = agent.__agent_id__ + self._steps = steps + self._agent_ids = agent_ids + self._executor = executor + self._step = 0 + + def _next_step(self) -> str: + step_id = f"step-{self._step:04d}" + self._step += 1 + return step_id + + @implements(step) + def _step(self, template: Callable, *args, **kwargs) -> Awaitable: + return self._run_step(self._next_step(), template, args, kwargs) + + @implements(scatter) + def _scatter_items(self, items, agent, fn) -> Awaitable: + return self._scatter(self._next_step(), items, agent, fn) + + def _step_within_item(self, template: Callable, *args, **kwargs) -> Awaitable: + """`step`, as interpreted while this agent runs a scatter item. + + The item already is a step, with its own ID and its own place in the + log, so there is nothing left to coordinate -- just run the call. + """ + agent = getattr(template, "__agent__", None) + if agent is not None and agent.__agent_id__ != self._agent_id: + raise RuntimeError( + f"a scatter item taken on by {self._agent_id!r} called " + f"{template.__name__}(), which belongs to " + f"{agent.__agent_id__!r}. A scatter item is work one agent " + f"does alone; step across agents outside the scatter." + ) + return self._in_thread(template, *args, **kwargs) + + async def _in_thread[T](self, fn: Callable[..., T], *args, **kwargs) -> T: + """Await *fn* on a worker thread, carrying the current context along. + + The context copy is what puts the agent's `effectful` handler stack -- + provider, retries, persistence -- in scope inside the worker. + """ + loop = asyncio.get_running_loop() + ctx = contextvars.copy_context() + return await loop.run_in_executor( + self._executor, functools.partial(ctx.run, fn, *args, **kwargs) + ) + + async def _run_step( + self, step_id: str, template: Callable, args: tuple, kwargs: dict + ) -> Any: + agent = getattr(template, "__agent__", None) + + if agent is None: + # Unbound template: not owned by anyone, so every agent runs it. + return await self._in_thread(template, *args, **kwargs) + + if self._agent_ids is not None and agent.__agent_id__ not in self._agent_ids: + raise ChoreographyError( + f"{template.__name__}() belongs to agent " + f"{agent.__agent_id__!r}, which is not part of this " + f"choreography -- no one would ever run it." + ) + + result = self._steps.result(step_id) + if agent.__agent_id__ != self._agent_id: + return await result + if result.done(): + # Recorded by an earlier run's log. + return result.result() + + try: + value = await self._in_thread(template, *args, **kwargs) + except Exception as e: + _fail(result, e) + raise + self._steps.resolve(step_id, value) + return value + + async def _scatter[A: Agent, T, U]( + self, + step_id: str, + items: Sequence[T], + agent: A | Sequence[A], + fn: Callable[[A, T], Awaitable[U]], + ) -> list[U]: + agents = [agent] if isinstance(agent, Agent) else list(agent) + results = [self._steps.result(f"{step_id}:{i}") for i in range(len(items))] + me = typing.cast(A, self._agent) + + if self._agent_id in {a.__agent_id__ for a in agents}: + work = self._steps.work(step_id, results) + while True: + try: + index = work.get_nowait() + except asyncio.QueueEmpty: + break + try: + # Rebinding `step` is what stops per-item work from + # allocating step IDs; the binding lasts exactly as long as + # the item does. + with handler({step: self._step_within_item}): + value = await fn(me, items[index]) + except Exception as e: + _fail(results[index], e) + raise + self._steps.resolve(f"{step_id}:{index}", value) + + return [await result for result in results] + + +# ── Choreography runner ─────────────────────────────────────────── + + +class Choreography[**P, T]: + """Run a choreographic program with endpoint projection. + + A `Choreography` is callable with the program's own signature, so it is + the program made runnable: where *program* is an ``async`` function + returning ``T``, the choreography is a plain callable returning ``T``, + having run every agent through it. + + Every agent runs *program* as its own `asyncio.Task`; `EndpointProjection` + is what makes each of those tasks behave differently. Blocking template + calls go to a thread pool sized to the number of agents, so no agent can be + starved by another's model call. + + The tasks run in an `asyncio.TaskGroup`, which supplies the parts the + threaded version had to build by hand: the first failure cancels the other + agents, and the failure propagates to the caller. Cancellation cannot + interrupt an LLM call that is already in flight on a worker thread, so a + failing run waits for those to return before it raises. + + Handlers are taken from the surrounding context, exactly as anywhere else + in `effectful`: install them with `~effectful.ops.semantics.handler` + around the run and every agent task inherits them, as does every worker + thread the agents call into. Nothing needs to be handed to the + choreography, which is also why a script run under + `effectful.handlers.llm.harness` needs no handler code of its own. + + Without a *log*, each run starts from a clean slate: results live in + memory for the duration of the run, so re-running a choreography + re-executes it. With one, each step is written to SQLite as it completes + and a later run replays what is already there, resuming at the first step + that never finished. Only successful steps are recorded, so a step that + failed or was interrupted simply runs again, and scatter items are + recorded one by one -- interrupt a scatter over ten modules after six and + the next run implements the remaining four. + + .. warning:: + + Steps are identified by position, so a log only makes sense for the + program that wrote it: editing the choreography shifts the IDs and the + recorded results land on the wrong steps. Delete the file, or use a + fresh path, whenever the program changes. + + Results are pickled, which is what lets a step return a dataclass or any + other decoded value rather than only JSON. A log is a cache of your own + run, read back with the same trust as + `~effectful.handlers.llm.completions.SQLitePersister`'s checkpoints -- and + read back only by running the choreography again, since a step ID means + nothing without the program that assigned it. + + Args: + program: The choreographic ``async`` function. All agents run it. + agents: The agents participating in the choreography. + log: Path to a SQLite database in which to record completed steps, so + that an interrupted run resumes when run again. ``None`` keeps + everything in memory. + + Example:: + + choreo = Choreography(build_codebase, agents=[architect, coder, reviewer]) + + with handler(LiteLLMProvider(model="gpt-4o-mini")), handler(RetryLLMHandler()): + result = choreo( + "Build a library...", + architect=architect, + coder=coder, + reviewer=reviewer, + ) + """ + + program: Callable[P, Awaitable[T]] + agents: list[Agent] + log: pathlib.Path | None + _steps: _Steps + + def __init__( + self, + program: Callable[P, Awaitable[T]], + agents: Sequence[Agent], + log: str | os.PathLike[str] | None = None, + ) -> None: + self.program = program + self.agents = list(agents) + self.log = pathlib.Path(log) if log is not None else None + self._steps = _Steps(self.log) + + async def run_async(self, *args: P.args, **kwargs: P.kwargs) -> T: + """Run the choreography to completion. + + The arguments are the program's own, forwarded to every agent. They all + compute the same result; that result is returned. + + Raises: + ChoreographyError: If any agent fails. + """ + # Fresh state per run: futures belong to the loop that created them. + self._steps = _Steps(self.log) + self._steps.replay() + agent_ids = frozenset(a.__agent_id__ for a in self.agents) + + async def as_agent(agent: Agent, executor: concurrent.futures.Executor) -> T: + projection = EndpointProjection( + agent, self._steps, agent_ids, executor=executor + ) + with handler(projection): + try: + return await self.program(*args, **kwargs) + except (asyncio.CancelledError, ChoreographyError): + raise + except Exception as e: + raise ChoreographyError( + f"Agent {agent.__agent_id__!r} failed: {e}" + ) from e + + tasks: list[asyncio.Task[T]] = [] + with concurrent.futures.ThreadPoolExecutor( + max_workers=max(1, len(self.agents)), thread_name_prefix="choreo" + ) as executor: + try: + async with asyncio.TaskGroup() as group: + tasks = [ + group.create_task( + as_agent(agent, executor), + name=f"choreo-{agent.__agent_id__}", + ) + for agent in self.agents + ] + except BaseExceptionGroup as group_error: + # Report one agent's failure rather than a group of one. The + # failure already names the agent; the group nests if the + # program runs task groups of its own. + error: BaseException = group_error + while isinstance(error, BaseExceptionGroup): + error = error.exceptions[0] + raise error + + return tasks[0].result() + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: + """Run the choreography from synchronous code. + + Equivalent to ``asyncio.run(choreo.run_async(...))``; await `run_async` + from inside an event loop that is already running. + """ + return asyncio.run(self.run_async(*args, **kwargs)) diff --git a/docs/source/llm_examples/choreographies/multi_agent_choreography.py b/docs/source/llm_examples/choreographies/multi_agent_choreography.py new file mode 100644 index 000000000..13a6564e3 --- /dev/null +++ b/docs/source/llm_examples/choreographies/multi_agent_choreography.py @@ -0,0 +1,376 @@ +"""Multi-agent library build via choreographic endpoint projection. + +Demonstrates: +- Choreographic programming: one ``async`` function describes the whole workflow +- Endpoint projection: every agent runs that function as its own `asyncio.Task`, + executing the steps it owns and awaiting the ones it doesn't +- ``scatter``: two coders share the implementation work and two reviewers share + the reviews, each item going to whichever agent is free +- A step log: interrupt the run and start it again, and the agents resume from + the last step that finished +- Tools as ground truth: the reviewers run each module's tests rather than + judging the code by reading it, so the fix loop turns on a fact + +The scenario: a team of agents collaboratively builds a small Python library. +An architect breaks the project into module specs, coders implement the modules +in parallel, and reviewers run their tests and review them in parallel, sending +work back to the coders until everything passes. + +The reviewers run generated test files as subprocesses, so this example +executes code the model wrote. Everything under +`effectful.handlers.llm.harness` already can -- it installs a Python REPL -- +but it is worth knowing before pointing this at an untrusted project spec. + +Only in-flight LLM calls occupy threads: an agent waiting on a peer's step is a +suspended coroutine. See ``library.py`` alongside this example for why steps are +spelled ``await step(...)`` rather than as plain method calls. + +Run it, interrupt it with Ctrl-C, and run it again to watch it pick up where it +left off:: + + python -m effectful.handlers.llm.harness \\ + docs/source/llm_examples/choreographies/multi_agent_choreography.py --model gpt-4o-mini + +Use ``--restart`` to forget the recorded steps and build from scratch, and pass +``--persist-db PATH`` to the harness to checkpoint each agent's own +conversation history alongside them. +""" + +import argparse +import json +import pathlib +import subprocess +import sys +from collections.abc import Sequence +from typing import Literal, TypedDict + +from docs.source.llm_examples.choreographies.library import ( + Choreography, + ChoreographyError, + scatter, + step, +) +from effectful.handlers.llm import Agent, Template, Tool + +DEFAULT_TEST_TIMEOUT = 60 +"""Seconds a generated test file gets before the reviewer gives up on it.""" + +# The project to build +PROJECT_SPEC = """\ +Build a small Python utility library called 'textkit' with these modules: +1. textkit/slugify.py — convert strings to URL-safe slugs +2. textkit/wrap.py — word-wrap text to a given width +3. textkit/redact.py — redact email addresses and phone numbers from text +Each module should have a clear public API, docstrings, and at least 3 +test cases written as a separate test_.py file. +""" + + +# --------------------------------------------------------------------------- +# Structured output — constrained decoding for LLM output +# --------------------------------------------------------------------------- + + +class ModuleSpec(TypedDict): + """Schema for architect planning output — constrained decoding ensures valid shape.""" + + module_path: str + description: str + public_api: str + test_path: str + + +class PlanResult(TypedDict): + """Wrapper for list output — LiteLLM requires a root object, not bare array.""" + + modules: list[ModuleSpec] + + +class ReviewResult(TypedDict): + """Schema for reviewer output — verdict constrained to PASS or NEEDS_FIXES.""" + + verdict: Literal["PASS", "NEEDS_FIXES"] + feedback: str + + +# --------------------------------------------------------------------------- +# Agents +# --------------------------------------------------------------------------- + + +class ArchitectAgent(Agent): + """You are a software architect. Given a project specification, you break + it into individual module implementation tasks. Each task should specify + the module filename, its public API, and what tests to write. + Be concrete and specific — the coder will follow your spec exactly. + """ + + def __init__(self, output_dir: pathlib.Path, **kwargs): + super().__init__(**kwargs) + self.output_dir = output_dir + + @Tool.define + def read_existing_files(self) -> str: + """List files already written to the output directory.""" + files = sorted(self.output_dir.rglob("*.py")) + if not files: + return "No Python files yet." + return "\n".join(str(f.relative_to(self.output_dir)) for f in files) + + @Template.define + def plan_modules(self, project_spec: str) -> PlanResult: + """Given this project specification, output a plan with a "modules" list. + Each module spec has: module_path, description, public_api, test_path. + + Use `read_existing_files` to check what's already been written + and skip those. + + Project spec: + {project_spec}""" + + +class CoderAgent(Agent): + """You are an expert Python developer. Given a module specification, + you write clean, well-documented Python code. You also write thorough + test files. Output ONLY the Python source code, no markdown fences. + """ + + def __init__(self, output_dir: pathlib.Path, **kwargs): + super().__init__(**kwargs) + self.output_dir = output_dir + + @Tool.define + def read_file(self, path: str) -> str: + """Read a file from the output directory.""" + full = self.output_dir / path + return full.read_text() if full.exists() else f"File not found: {path}" + + @Tool.define + def write_file(self, path: str, content: str) -> str: + """Write a file to the output directory.""" + full = self.output_dir / path + full.parent.mkdir(parents=True, exist_ok=True) + full.write_text(content) + return f"Wrote {len(content)} chars to {path}" + + @Template.define + def implement_module(self, module_spec: str) -> str: + """Implement the following module specification. Use `write_file` + to write both the module and its test file. Use `read_file` to + check existing code if needed. + + Specification: + {module_spec}""" + + +class ReviewerAgent(Agent): + """You are a senior code reviewer. You review Python modules for + correctness, style, edge cases, and test coverage. You judge a module by + running its tests, not only by reading it. Be specific about issues and + provide actionable feedback. + """ + + def __init__( + self, + output_dir: pathlib.Path, + test_timeout: float = DEFAULT_TEST_TIMEOUT, + **kwargs, + ): + super().__init__(**kwargs) + self.output_dir = output_dir + self.test_timeout = test_timeout + + @Tool.define + def read_file(self, path: str) -> str: + """Read a file from the output directory.""" + full = self.output_dir / path + return full.read_text() if full.exists() else f"File not found: {path}" + + @Tool.define + def run_tests(self, test_path: str) -> str: + """Run the test file at `test_path` with pytest and return its output.""" + try: + result = subprocess.run( + # `-o addopts=` because pytest would otherwise inherit the + # addopts of whatever project the workspace happens to sit in. + [sys.executable, "-m", "pytest", test_path, "-q", "--no-header"] + + ["-o", "addopts=", "-p", "no:cacheprovider"], + cwd=self.output_dir, + capture_output=True, + text=True, + timeout=self.test_timeout, + ) + except subprocess.TimeoutExpired: + return f"Timed out after {self.test_timeout}s — the tests do not terminate." + return f"exit code {result.returncode}\n\n{result.stdout[-4000:]}" + + @Template.define + def review_module(self, module_path: str, test_path: str) -> ReviewResult: + """Review the module at {module_path} and its tests at {test_path}. + Use `read_file` to read them and `run_tests` to run the test file. + + Return verdict "PASS" or "NEEDS_FIXES" and feedback. A module whose + tests do not all pass is "NEEDS_FIXES", whatever the code looks like; + say which test failed and why. If the test itself is wrong, say that + instead — either way the coder has something to fix.""" + + +# --------------------------------------------------------------------------- +# Choreographic program — the entire multi-agent workflow in one function +# --------------------------------------------------------------------------- + + +async def build_project( + project_spec: str, + architect: ArchitectAgent, + coder: CoderAgent | Sequence[CoderAgent], + reviewer: ReviewerAgent | Sequence[ReviewerAgent], + max_rounds: int, +) -> list[ReviewResult]: + """Choreographic program describing the full build workflow. + + A role may be filled by one agent or by several: `scatter` hands each item + to whichever of them is free, which is why the coder and reviewer + parameters are typed to accept a pool. + + 1. Architect breaks the project into module specs. + 2. Coders implement modules in parallel (scatter hands each to whoever is free). + 3. Reviewers run each module's tests and review it; coders fix what failed, + for up to *max_rounds* rounds. + """ + # Step 1: the architect plans the modules. Every agent awaits this same + # step; only the architect calls the model for it. + plan = await step(architect.plan_modules, project_spec) + + # Step 2: scatter implementation across the coders. Each coder takes the + # next module as it becomes free, until none are left. + await scatter( + plan["modules"], + coder, + lambda c, mod: step(c.implement_module, json.dumps(mod, indent=2)), + ) + + # Step 3: review loop — keep fixing until the reviewers accept every module. + # Bounded, because a reviewer and a coder that disagree would otherwise + # trade rounds forever. Every agent sees the same reviews, so they all + # leave the loop on the same iteration. + for _ in range(max_rounds): + reviews: list[ReviewResult] = await scatter( + plan["modules"], + reviewer, + lambda r, mod: step(r.review_module, mod["module_path"], mod["test_path"]), + ) + + needs_fixes = [ + (mod, review) + for mod, review in zip(plan["modules"], reviews) + if review["verdict"] == "NEEDS_FIXES" + ] + if not needs_fixes: + return reviews + + await scatter( + needs_fixes, + coder, + lambda c, pair: step( + c.implement_module, + json.dumps({**pair[0], "fix_feedback": pair[1]["feedback"]}, indent=2), + ), + ) + + return reviews # out of rounds; hand back the last verdicts as they stand + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--workspace", + type=pathlib.Path, + default=pathlib.Path("./multi_agent_workspace"), + help="Directory to write the generated library into", + ) + parser.add_argument( + "--project-spec", + type=str, + default=PROJECT_SPEC, + help="The project for the team to build", + ) + parser.add_argument("--coders", type=int, default=2, help="Number of coder agents") + parser.add_argument( + "--reviewers", type=int, default=2, help="Number of reviewer agents" + ) + parser.add_argument( + "--max-rounds", + type=int, + default=3, + help="How many review-and-fix rounds to allow before giving up", + ) + parser.add_argument( + "--test-timeout", + type=float, + default=DEFAULT_TEST_TIMEOUT, + metavar="SECONDS", + help="How long a reviewer waits for a generated test file to finish", + ) + parser.add_argument( + "--restart", + action="store_true", + help="Forget the steps recorded by earlier runs and build from scratch", + ) + args = parser.parse_args() + + output_dir = args.workspace / "output" + output_dir.mkdir(parents=True, exist_ok=True) + + # An explicit agent_id is what makes an Agent persistent, and it is also how + # endpoint projection tells the agents apart. + architect = ArchitectAgent(output_dir, agent_id="architect") + coders = [CoderAgent(output_dir, agent_id=f"coder-{i}") for i in range(args.coders)] + reviewers = [ + ReviewerAgent(output_dir, args.test_timeout, agent_id=f"reviewer-{i}") + for i in range(args.reviewers) + ] + + # Steps completed by an earlier run are replayed instead of re-asking the + # model, so an interrupted build resumes rather than starting over. + log = args.workspace / ".state" / "steps.db" + if args.restart: + log.unlink(missing_ok=True) + # Ask before building the choreography, which creates the log if it is new. + resuming = log.exists() + + # Tasks, the thread pool and cancellation on failure are all handled for + # you; the model handlers come from the harness. + choreo = Choreography( + build_project, agents=[architect, *coders, *reviewers], log=log + ) + + print(f"{'Resuming' if resuming else 'Starting'} multi-agent build") + try: + reviews = choreo( + args.project_spec, + architect=architect, + coder=coders, + reviewer=reviewers, + max_rounds=args.max_rounds, + ) + except ChoreographyError as e: + print(f"Choreography failed: {e} — re-run to retry from this step") + return + except KeyboardInterrupt: + print("Interrupted — re-run to resume from the last completed step") + return + + passed = sum(1 for r in reviews if r["verdict"] == "PASS") + print(f"\nDone: {len(reviews)} modules reviewed, {passed} passed") + for f in sorted(output_dir.rglob("*.py")): + print(f" {f.relative_to(args.workspace)}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/__init__.py b/docs/source/llm_examples/optimization/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/optimization/kernels.py b/docs/source/llm_examples/optimization/kernels.py new file mode 100644 index 000000000..fbcbabdef --- /dev/null +++ b/docs/source/llm_examples/optimization/kernels.py @@ -0,0 +1,798 @@ +"""Kernel instructions: multi-task search over a shared frontier (optimize_anything 5.2). + +A dataset of related problems is supplied but no validation set, which selects the +paper's multi-task mode -- the one no prior LLM-evolution framework has. The frontier is +shared across tasks so a pattern discovered while working on one is available as a +parent when proposing for another, and at output time each task independently picks its +own best candidate off that frontier. Multi-task search therefore produces N specialized +artifacts that have all benefited from a common search, which is the distinction the +paper draws against generalization mode's single artifact. + +The artifact is the *instruction that drives code generation*, exactly as the paper +evolves the prompt behind its CUDA kernels rather than the kernels themselves. Each +evaluation hands that instruction to a cheaper programmer model (``--worker-model``), +which writes the function, and scores what comes back. + +Scoring borrows KernelBench's *shape* -- correctness against a reference +implementation, then wall-clock speedup against it -- and that shape is what makes the +domain optimizable at all: the worker nearly always writes a correct list transform from +the bare seed instruction -- 14 of the 15 seed evaluations across the runs below -- so +correctness alone would saturate immediately. Correctness is not, however, a floor the +search stays above. An instruction that pushes hard for speed makes the worker write +kernels that fail, and those score zero; the gate is a cliff the search repeatedly falls +off, as the traces below record. + +Read the speedups with the baseline in mind, because it is not the paper's. KernelBench +compares against PyTorch, which is cuDNN and cuBLAS -- vendor-tuned code, and the reason +"87% match or beat the baseline" is a strong claim. The reference implementations here +are deliberately plain Python, explicit loops and ``append``, slow enough that ruff +objects to them in as many words. Beating them by 1.6x is beating unoptimized +interpreter code, not a tuned library, and the two numbers are not comparable. + +The five tasks share their *failure modes* rather than their algorithm, which is what +gives cross-transfer something to transfer: three have a naive formulation that rescans +the whole prefix and is quadratic, and two turn on degenerate inputs a hurried +implementation skips. Both lessons are worth less than they sound. The references +already carry running state, so an instruction that teaches it recovers the baseline +rather than beating it -- 0.0 to about 1.0, a timeout fix wearing a speedup's clothes -- +and the degenerate cases are stated in the specification text the worker is handed, so +the transferable insight there is "read the specification". + +Demonstrates: +- Multi-task mode: per-task Pareto objectives on one shared frontier, per-task winners + at output time, and a count of how many of those winners were last refined while the + proposer was looking at a *different* task -- reported against the count chance alone + would produce, because with five tasks and a two-task minibatch that null covers most + of the statistic's range and the raw count says nothing on its own +- A single-task control (``--single-task``) that re-optimizes each task independently, + which is the comparison the paper's 5.4 reports -- though see the simplifications on + what "equivalent budget" does and does not mean here +- Side Information as compiler-style feedback: failing cases with expected and actual + values, measured times and speedup, the traceback, and the code itself +- A correctness gate that a search for speed can and does fall foul of + +Measured on 2026-07-30 with gpt-5.5 proposing and gpt-4.1-mini writing kernels. The +score is the mean over the five per-task winners of their speedup against the reference +implementation: + + multi-task, 10 iterations 1.385 -> 1.569 (40 evaluator calls) + best single artifact 1.385 (the matched comparison) + single-task control, 2 iterations/task 1.328 -> 1.396 (15 evaluator calls) + single-task control, 7 iterations/task 0.999 -> 1.638 (40 evaluator calls) + +The paper's 5.4 finding -- multi-task ahead of single-task at equivalent per-problem +budget -- comes out whichever way the budget is counted. Matched on optimizer +iterations, the multi-task arm leads, 1.569 against 1.396, and spends 2.7x the evaluator +calls doing it. Matched on evaluator calls instead, the control leads, 1.638 against +1.569, and is now the profligate arm on the other currency: 35 proposer calls against +10. No setting of the two knobs matches both, because one multi-task iteration buys a +five-task evaluation and one single-task iteration buys one, while both buy exactly one +proposer call. + +Nothing here separates those arms from noise. The seed is the same string in all three +runs and scored 1.385, 1.328 and 0.999 on the same five tasks -- a 39% spread, wider +than any gap between the arms -- because the worker rewrites the kernel from scratch +every time, and on one of the three draws never produced a usable ``zscore`` kernel at +all. Per task the spread is worse: ``window_sum_101``'s seed scored 0.762 and 0.988, +``zscore``'s 2.114, 1.642 and 0.000. + +The multi-task headline is also a per-task maximum over a six-candidate pool, and the +matched one-artifact-against-one-artifact number the report prints next to it says what +that is worth here: the best single artifact scored 1.385, which is the seed's own mean. +No instruction the search wrote beat the instruction it started from, averaged over the +five tasks. The whole of that arm's gain is composition -- four different candidates, +each best at one or two tasks. + +The correctness cliff is visible in the traces rather than in the summary numbers. The +multi-task arm proposed an instruction that turned a 2.02 minibatch into a 0, and the +evaluation-matched control's ``l2_normalize`` run scored exactly zero on four of its +seven proposals: an instruction pushing hard enough for speed makes the worker write +kernels that are wrong, and wrong scores nothing however fast it is. + +Cross-task transfer came out at 3 of the 4 refined winners last refined while the +proposer was looking at a different task, against the 2.4 chance alone would give. Being +0.6 of a winner above chance on a statistic that can take five values is not evidence of +anything. The fifth winner was the unrefined seed, and is excluded from both figures. + +What this domain demonstrates is the machinery -- a shared frontier, per-task selection +off it, a control to compare against, and a transfer statistic reported against its null +-- not a result. One run of each arm on five tasks, against the paper's 31, through a +noise floor that swallows the effect, decides nothing in either direction. + +One thing any number here includes and cannot be separated from: the harness's +``TenacityRetryer`` sits above the worker model, so a kernel whose source does not +decode is fed its own error and asked again. Every speedup is therefore a speedup for +the instruction *plus that repair loop*, and an instruction that provokes +borderline-undecodable code is flattered by it. +""" + +# Simplifications vs. the source: +# - Pure-Python list transforms on a CPU, not CUDA kernels on a V100 against +# KernelBench's 31 PyTorch operations, and no NVCC in the loop. The baseline is +# plain Python rather than a vendor-tuned library, so a speedup here is not the same +# quantity the paper reports (see the header). +# - The score is mean speedup rather than the paper's fast_p(s) curve, and the side +# information has no documentation-retrieval channel. With five tasks, fast_p could +# be reported in 20% increments from the same per-task scores; it is not. +# - Budget is counted in optimizer iterations, not metric calls or dollars; the paper +# spends ~3000 metric calls and $140 on this domain. +# - There is no budget in which the two arms are matched. Matching iterations, as +# ``--single-task`` does by default, leaves multi-task with 2.7x the evaluator calls +# (40 against 15) -- it pays for an evaluation across all five tasks whenever a +# proposal is accepted, while the control pays for one. Matching evaluator calls with +# ``--control-iterations 7`` inverts it: the control then gets 35 proposals to +# multi-task's 10, and reflection is the expensive call. ``report`` prints both counts +# for both arms so which currency a comparison is stated in stays visible. Multi-task +# also takes its per-task maximum over a larger pool, which favours it for reasons +# unrelated to transfer. +# - Nor can the two arms be made to differ in exactly one way. Single-task mode has no +# per-example objectives to keep a frontier over, so choosing it necessarily changes +# both the task count and what the Pareto objectives are. That is why the paper +# introduces per-metric objectives for that mode, and it is a property of its own +# comparison as much as of this one. +# - The headline gain is upward-biased on the multi-task side and cannot be negative +# there: the result is a per-task maximum over the whole pool while the seed is a +# single candidate's mean, so per-task max is >= the seed's score on every task by +# construction. The report prints the best *single* artifact's mean alongside it, +# which is the matched one-artifact-to-one-artifact comparison. +# - One run per arm, and no variance estimate beyond the seed. That seed is measured +# afresh in every run, which is the only repeated measurement here and is enough to +# settle the question: the same instruction on the same five tasks spans 0.999 to +# 1.385, so the noise floor is several times the effects being compared. Most of that +# is not timing jitter -- the worker resamples the kernel each run, so a task's seed +# score can move by 30% or drop to zero on a synthesis that never decodes. +# - Cross-task transfer is a lineage count over one run against one control, not the +# paper's MT10/MT20 scaling study. It inspects only the last refinement rather than +# full ancestry, so it is reported against its null expectation rather than alone. +# - The score is a wall-clock ratio, so it is only as reproducible as the machine is +# quiet. The baseline is re-timed back to back with every candidate for exactly this +# reason (see `measure_speedup`), which makes the ratio robust to a loaded machine but +# not to one whose speed changes mid-measurement. + +import argparse +import bisect +import collections.abc +import math +import random +import signal +import statistics +import threading +import time +import traceback +import zlib + +import pydantic.dataclasses + +from docs.source.llm_examples.optimization.library import ( + WORKER_MODEL, + Candidate, + Diagnostic, + Evaluation, + Metric, + Result, + Rollout, + optimize_anything, + report, + source_of, + worker, +) +from effectful.handlers.llm import Agent, Template + +# A kernel is one list-to-list transform; every task in the family shares this +# signature so a single instruction can drive all of them. +type Kernel = collections.abc.Callable[[list[float]], list[float]] + + +class _Timeout(Exception): + pass + + +@pydantic.dataclasses.dataclass(frozen=True) +class KernelTask: + """One task in the family: what the transform must compute. The test cases are + hidden -- they live in ``KERNEL_TESTS``, and reach the proposer only as the + specific failures reported in Side Information.""" + + name: str + spec: str + + +# A family that shares its *failure modes* rather than its algorithm, which is what +# makes cross-transfer possible at all. Two lessons run through it. Three of the tasks +# have a naive formulation that recomputes over the whole prefix or window and is +# quadratic -- correct on the small cases, far too slow on the timed one -- so the +# transferable lesson is "carry running state instead of rescanning". The other two +# turn on degenerate inputs the specification states and a hurried implementation +# skips. An instruction that learns either lesson on one task collects points on the +# others, exactly as the paper's CUDA instruction learns coalescing once and spends it +# across 31 kernels. +KERNEL_TASKS: list[KernelTask] = [ + KernelTask( + "count_smaller_before", + "For each position i, output the number of earlier positions j < i whose value " + "is strictly smaller than the value at i. The output has the same length as " + "the input; an empty input gives an empty output.", + ), + KernelTask( + "window_sum_101", + "For each position i, output the sum of the values from index max(0, i - 100) " + "through i inclusive -- a trailing window of up to 101 values, shorter near " + "the start. The output has the same length as the input; an empty input gives " + "an empty output.", + ), + KernelTask( + "distinct_prefix_counts", + "For each position i, output how many distinct values occur in the input up to " + "and including position i. The output has the same length as the input; an " + "empty input gives an empty output.", + ), + KernelTask( + "l2_normalize", + "Divide every value by the Euclidean (L2) norm of the whole input, so the " + "result has unit norm. If that norm is exactly zero, every output value is " + "0.0. The output has the same length as the input; an empty input gives an " + "empty output.", + ), + KernelTask( + "zscore", + "Standardize the input: subtract the mean and divide by the population " + "standard deviation. If that standard deviation is exactly zero, every output " + "value is 0.0. The output has the same length as the input; an empty input " + "gives an empty output.", + ), +] + +type Case = tuple[list[float], list[float]] + +# The small cases are the contract, written out so the edge semantics are readable +# rather than implied by a reference implementation. +KERNEL_TESTS: dict[str, list[Case]] = { + "count_smaller_before": [ + ([], []), + ([5.0], [0.0]), + ([3.0, 1.0, 2.0], [0.0, 0.0, 1.0]), + ([2.0, 2.0, 1.0, 4.0], [0.0, 0.0, 0.0, 3.0]), + ], + "window_sum_101": [ + ([], []), + ([5.0], [5.0]), + ([1.0, 2.0, 3.0], [1.0, 3.0, 6.0]), + ([-1.0, 1.0, -1.0, 1.0], [-1.0, 0.0, -1.0, 0.0]), + ], + "distinct_prefix_counts": [ + ([], []), + ([5.0], [1.0]), + ([1.0, 1.0, 2.0], [1.0, 1.0, 2.0]), + ([3.0, 1.0, 3.0, 2.0], [1.0, 2.0, 2.0, 3.0]), + ], + "l2_normalize": [ + ([], []), + ([0.0, 0.0], [0.0, 0.0]), + ([3.0, 4.0], [0.6, 0.8]), + ([-3.0, 4.0], [-0.6, 0.8]), + ], + "zscore": [ + ([], []), + ([5.0, 5.0, 5.0], [0.0, 0.0, 0.0]), + ([2.0], [0.0]), + ([1.0, 2.0, 3.0], [-1.224744871391589, 0.0, 1.224744871391589]), + ], +} + + +# Written the plain way on purpose: explicit loops, ``append`` per element, arithmetic +# spelled out. This is the "straightforward implementation" a competent programmer +# reaches for first, and it is the baseline the score is a ratio against -- the role +# KernelBench's unoptimized PyTorch reference plays in the paper. Rewriting these with +# comprehensions, ``itertools.accumulate``, locally bound methods or a reciprocal +# multiply is exactly the headroom the search is asked to find. (The ``noqa``s below +# are load-bearing: ruff is right that a comprehension would be faster, and being +# slower than that is precisely this code's job.) + + +def _count_smaller_before(values: list[float]) -> list[float]: + counts: list[float] = [] + seen: list[float] = [] + for x in values: + counts.append(float(bisect.bisect_left(seen, x))) + bisect.insort(seen, x) + return counts + + +def _window_sum_101(values: list[float]) -> list[float]: + out: list[float] = [] + running = 0.0 + for i in range(len(values)): + running = running + values[i] + if i >= 101: + running = running - values[i - 101] + out.append(running) + return out + + +def _distinct_prefix_counts(values: list[float]) -> list[float]: + out: list[float] = [] + seen: set[float] = set() + for i in range(len(values)): + seen.add(values[i]) + out.append(float(len(seen))) + return out + + +def _l2_normalize(values: list[float]) -> list[float]: + total = 0.0 + for i in range(len(values)): + total = total + values[i] * values[i] + norm = math.sqrt(total) + out: list[float] = [] + for i in range(len(values)): + out.append(0.0 if norm == 0.0 else values[i] / norm) # noqa: PERF401 + return out + + +def _zscore(values: list[float]) -> list[float]: + if not values: + return [] + total = 0.0 + for i in range(len(values)): + total = total + values[i] + mean = total / len(values) + variance = 0.0 + for i in range(len(values)): + variance = variance + (values[i] - mean) * (values[i] - mean) + deviation = math.sqrt(variance / len(values)) + out: list[float] = [] + for i in range(len(values)): + out.append( # noqa: PERF401 + 0.0 if deviation == 0.0 else (values[i] - mean) / deviation + ) + return out + + +# The baseline every candidate is measured against -- the straightforward linear +# implementation a competent programmer writes without thinking about speed. It is +# never shown to the model; it supplies the expected output of the timed case and the +# denominator of the speedup, exactly as KernelBench's PyTorch baseline does in the +# paper's 5.2. +REFERENCE: dict[str, Kernel] = { + "count_smaller_before": _count_smaller_before, + "window_sum_101": _window_sum_101, + "distinct_prefix_counts": _distinct_prefix_counts, + "l2_normalize": _l2_normalize, + "zscore": _zscore, +} + +# Size of the timed input per task, and the wall-clock ceiling that stops a quadratic +# implementation instead of letting it hang the run. The sizes are chosen so the +# reference finishes in tens of milliseconds and a rescanning implementation cannot +# finish at all: 30k elements of prefix rescanning is ~450M comparisons. +KERNEL_PERF: dict[str, tuple[int, float]] = { + "count_smaller_before": (30_000, 5.0), + "window_sum_101": (300_000, 5.0), + "distinct_prefix_counts": (300_000, 5.0), + "l2_normalize": (300_000, 5.0), + "zscore": (300_000, 5.0), +} + +TIMING_REPEATS = 3 # best of three: the standard robust estimator for a short run + + +def perf_input(task: str) -> list[float]: + """The timed case's input: deterministic pseudo-random values, so every candidate + is timed on exactly the same work. Seeded from a checksum of the task name rather + than its length, which silently gave two tasks the same input the moment their + names happened to match in length.""" + rng = random.Random(zlib.crc32(task.encode())) + return [rng.uniform(-1.0, 1.0) for _ in range(KERNEL_PERF[task][0])] + + +def check_reference_agrees() -> bool: + """The reference implementations must reproduce the written-out contract, or the + timed case would be testing a different function than the small cases do. + + >>> check_reference_agrees() + True + """ + for name, cases in KERNEL_TESTS.items(): + for values, expected in cases: + produced = REFERENCE[name](list(values)) + assert len(produced) == len(expected) and all( + math.isclose(p, e, rel_tol=1e-9, abs_tol=1e-9) + for p, e in zip(produced, expected) + ), f"reference for {name} disagrees with the contract on {values}" + return True + + +SEED_INSTRUCTION = "Write a Python function that implements the specification." + + +class Programmer(Agent): + """You are an expert Python programmer. You implement exactly the specification + you are given, following the engineering instruction you are handed, and you + answer with code rather than prose.""" + + @Template.define + def write_kernel(self, instruction: str, task: KernelTask) -> Kernel: + """Write ``kernel(values)``: a function taking a ``list[float]`` and returning + a ``list[float]``, implementing this specification exactly. + + + {task.spec} + + + Follow this engineering instruction while you write it: + + + {instruction} + + + Standard library only. It is checked against hidden cases and then TIMED on a + large input: a correct implementation scores the ratio of a straightforward + reference implementation's time to yours, and an incorrect one scores zero + however fast it is. Write it to be both right and fast. + """ + + +def _time_kernel( + kernel: Kernel, values: list[float], ceiling: float +) -> tuple[list[float], float]: + """Best-of-``TIMING_REPEATS`` wall-clock time for one kernel on one input, with an + alarm so a quadratic implementation is stopped rather than left to hang.""" + + def _alarm(signum: int, frame: object) -> None: + raise _Timeout(f"exceeded the {ceiling}s ceiling on {len(values)} values") + + guarded = threading.current_thread() is threading.main_thread() + best, produced = math.inf, [] + if guarded: + previous = signal.signal(signal.SIGALRM, _alarm) + try: + for _ in range(TIMING_REPEATS): + if guarded: + signal.setitimer(signal.ITIMER_REAL, ceiling) + start = time.perf_counter() + produced = list(kernel(list(values))) + best = min(best, time.perf_counter() - start) + if guarded: + signal.setitimer(signal.ITIMER_REAL, 0.0) + finally: + if guarded: + signal.setitimer(signal.ITIMER_REAL, 0.0) + signal.signal(signal.SIGALRM, previous) + return produced, best + + +def measure_speedup(kernel: Kernel, task: str) -> tuple[float, Diagnostic]: + """Correctness-gated speedup over the reference on the large input. + + This is the paper's KernelBench metric in miniature: a kernel that is wrong scores + nothing, and a kernel that is right scores how many times faster than the baseline + it runs. It is also what keeps this domain from saturating -- every model writes a + correct list transform on the first try, so correctness alone would have nothing + left to optimize. It is not thereby a *floor* the search stays above: an instruction + that pushes hard for speed makes the worker write kernels that fail, and the run + logs show the search falling off that cliff repeatedly. + + The baseline is re-timed next to every candidate rather than measured once and + cached. That looks wasteful and is not: a wall-clock *ratio* is only meaningful if + both sides saw the same machine, and timing the reference on an idle process while + candidates are timed under load produces scores that swing by 5x with nothing about + the code having changed. Interleaving the two costs milliseconds and makes the + number reproducible. + """ + values, (size, ceiling) = perf_input(task), KERNEL_PERF[task] + expected, baseline = _time_kernel(REFERENCE[task], values, ceiling) + try: + produced, seconds = _time_kernel(kernel, values, ceiling) + except _Timeout as exc: + # Report the ceiling being hit and nothing else. Naming the likely cause here -- + # "rescanning earlier values for every position is quadratic; carry running + # state" -- would hand the proposer the lesson the search is then credited with + # discovering, and would be wrong besides on every other way of exceeding a + # ceiling. Diagnosing is the proposer's job; the evaluator's is to say + # accurately what happened. + return 0.0, Diagnostic("speed", f"on {size} values: {exc}") + except Exception as exc: + return 0.0, Diagnostic( + "speed", f"on {size} values this raised {type(exc).__name__}: {exc}" + ) + if len(produced) != len(expected) or not all( + math.isclose(p, e, rel_tol=1e-7, abs_tol=1e-7) + for p, e in zip(produced, expected) + ): + return 0.0, Diagnostic( + "speed", f"wrong output on the {size}-value input, so speed does not count" + ) + return baseline / seconds, Diagnostic( + "speed", + f"{size} values in {seconds * 1e3:.1f}ms against the reference " + f"implementation's {baseline * 1e3:.1f}ms measured back to back -- " + f"{baseline / seconds:.2f}x", + ) + + +def evaluate_instruction( + instruction: str, task: KernelTask | None, model: str +) -> Evaluation: + """Synthesize a kernel under the candidate instruction, check it, and time it. + + The score is the measured speedup over the reference implementation, gated on + correctness: any failing case scores zero, however fast the code is. That is the + paper's KernelBench setup (correctness against the reference, then wall-clock + against the PyTorch baseline), and it is what gives this domain something to climb. + The Side Information is the failing cases with expected and actual values, the + measured times and speedup, the traceback if it crashed, and the code itself. + + The metrics are only read in single-task mode, where the engine's Pareto objectives + are an evaluation's sub-scores rather than a dataset's examples -- which is what the + ``--single-task`` control runs. ``score`` has to be one of them because it is the + number the report reads as the headline. + """ + assert task is not None, "the kernel domain always has a dataset" + try: + with worker(model): + kernel = Programmer().write_kernel(instruction, task) + except Exception: + return Evaluation( + score=0.0, + diagnostics=[ + Diagnostic("task", f"{task.name}: {task.spec}"), + Diagnostic("synthesis failed", traceback.format_exc(limit=2).strip()), + ], + ) + + cases = KERNEL_TESTS[task.name] + passed = 0 + missed: list[Diagnostic] = [] + start = time.perf_counter() + for values, expected in cases: + try: + produced = list(kernel(list(values))) + ok = len(produced) == len(expected) and all( + math.isclose(p, e, rel_tol=1e-9, abs_tol=1e-9) + for p, e in zip(produced, expected) + ) + except Exception as exc: + ok, produced = False, f"raised {type(exc).__name__}: {exc}" # type: ignore[assignment] + if ok: + passed += 1 + else: + missed.append( + Diagnostic( + "failing case", + f"kernel({values}) returned {produced}, expected {expected}", + ) + ) + elapsed = time.perf_counter() - start + + speedup, timing = measure_speedup(kernel, task.name) + correct = passed == len(cases) and speedup > 0.0 + + # Only the first couple of failures go back: a wall of them buries the signal. Say + # how many were withheld, so the proposer is not told a partial list is the whole one. + failures = missed[:2] + if len(missed) > len(failures): + failures.append( + Diagnostic( + "further failures", + f"{len(missed) - len(failures)} more case(s) also failed and are not " + f"shown here", + ) + ) + diagnostics = [Diagnostic("task", f"{task.name}: {task.spec}")] + diagnostics += failures or [Diagnostic("correctness", "all small cases passed")] + diagnostics.append(timing) + diagnostics.append( + Diagnostic("small-case timing", f"{len(cases)} cases in {elapsed * 1e3:.2f}ms") + ) + diagnostics.append( + Diagnostic("code under test", (source_of(kernel) or "(unavailable)").strip()) + ) + diagnostics.append( + Diagnostic( + "verdict", + f"correct, and {speedup:.2f}x the reference implementation's speed -- the " + f"score IS that ratio, so a correct but ordinary implementation scores " + f"about 1.0 and only a faster one improves" + if correct + else f"{passed}/{len(cases)} small cases passed; an incorrect kernel " + f"scores zero no matter how fast it is", + ) + ) + return Evaluation( + score=speedup if correct else 0.0, + metrics=[ + Metric("score", speedup if correct else 0.0), + Metric("cases_passed", float(passed)), + ], + diagnostics=diagnostics, + ) + + +class Proposer(Agent): + """You are a reflective optimizer. You are shown the current instruction, the + scores the code written under it achieved, and diagnostic side information + explaining *why*, and you return a better instruction. You do not mutate blindly: + you first read the diagnostics to decide which failure mode is costing the most, + then you write the guidance that addresses it.""" + + @Template.define + def propose_instruction(self, current: str, feedback: list[Rollout]) -> str: + """You are optimizing the INSTRUCTION handed to a programmer model that + implements small list-transform functions. The instruction below is the + artifact -- it is reused for every task in a family, so it must say things + that are true of all of them. + + + {current} + + + Here is how the code written under it fared on a couple of tasks, including + the specific test cases that failed: + + + {feedback} + + + Diagnose the failures, then rewrite the instruction so a programmer following + it would not make them again. Prefer guidance that would still apply to a task + you have not been shown over anything specific to one task -- an instruction + that solves one task by naming its answer is worthless on the others. + + Return the improved instruction as plain text, nothing else. + """ + + +# --------------------------------------------------------------------------- +# Wiring and main +# --------------------------------------------------------------------------- + + +def run_kernel(args: argparse.Namespace, rng: random.Random) -> Result: + """Multi-task by default. ``--single-task`` runs the paper's control instead: each + task optimized independently, which is the comparison its 5.4 ablation reports. + + The control runs the engine's *single-task* mode, with the task bound in the + evaluator's closure and no dataset at all, so its Pareto objectives are the + evaluation's own sub-scores. Passing ``dataset=[task]`` would look equivalent and is + not: that is multi-task mode with one example, a frontier over a single objective on + which every tie is non-dominated and selection collapses to greedy, so the control + would differ from the treatment arm in its selection rule as well as its task count. + + It is worth being clear that there is still no configuration in which the two arms + differ by exactly one thing. Single-task mode necessarily changes both the number of + tasks *and* what the objectives are, since with one task there are no per-example + objectives to keep a frontier over -- which is precisely why the paper introduces + per-metric objectives for that mode. The paper's comparison has the same property. + + The two arms are also matched on optimizer iterations, not on evaluator calls, and + those are not the same thing -- multi-task pays for a full five-task evaluation + whenever a proposal is accepted. ``--control-iterations`` sets the control's per-task + budget directly, which is how to match on the evaluation counts ``report`` prints + for both arms. It buys that match with a mismatch elsewhere: raising the control's + per-task iterations raises its proposer calls in step, so an evaluation-matched + control makes several times as many reflection calls as the treatment arm. Both + counts are printed because neither currency can be held fixed alone. + """ + proposer = lambda instruction, feedback: Proposer().propose_instruction( # noqa: E731 + instruction, feedback + ) + if not args.single_task: + return optimize_anything( + evaluator=lambda i, t: evaluate_instruction(i, t, args.worker_model), + proposer=proposer, + seed=SEED_INSTRUCTION, + dataset=KERNEL_TASKS, + budget=args.budget, + minibatch_size=args.minibatch, + selection=args.selection, + use_side_info=not args.no_side_info, + rng=rng, + ) + + def evaluator_for(task: KernelTask) -> collections.abc.Callable[..., Evaluation]: + """Bind the task, so the engine sees a single-task problem with no dataset.""" + + def evaluate(instruction: str, _: object) -> Evaluation: + return evaluate_instruction(instruction, task, args.worker_model) + + return evaluate + + # The control is five independent runs, and its report has to be an aggregate of all + # five: reusing one run's ``Result`` and overwriting a few of its fields would print + # that run's frontier, iteration count and "best artifact" as though they were the + # whole control's. + runs: list[Result] = [] + per_task: list[tuple[str, Candidate, float]] = [] + per_problem = args.control_iterations or max(1, args.budget // len(KERNEL_TASKS)) + for task in KERNEL_TASKS: + print(f"\n[single-task control] {task.name} ({per_problem} iterations)") + result: Result = optimize_anything( + evaluator=evaluator_for(task), + proposer=proposer, + seed=SEED_INSTRUCTION, + budget=per_problem, + selection=args.selection, + use_side_info=not args.no_side_info, + rng=rng, + task_name=task.name, + ) + runs.append(result) + per_task.append((task.name, result.best, result.best_score)) + + best_run = max(runs, key=lambda r: r.best_score) + return Result( + mode=f"single-task control ({len(runs)} independent runs)", + # No pool and no objectives: this arm has five separate frontiers over five + # disjoint objective sets, and there is no honest way to merge them. Pooling + # the candidates would ask the Pareto machinery to compare a count_smaller_before + # score against a zscore one -- it raises, and it should. An empty pool tells + # `report` there is no shared frontier here, which is exactly the difference + # from the multi-task arm that this control exists to isolate. + pool=[], + history=[step for r in runs for step in r.history], + objectives=[], + seed_score=statistics.fmean(r.seed_score for r in runs), + best=best_run.best, + best_score=statistics.fmean(score for _, _, score in per_task), + per_task=per_task, + evaluations=sum(r.evaluations for r in runs), + proposals=sum(r.proposals for r in runs), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--budget", type=int, default=10, help="Optimizer iterations") + parser.add_argument( + "--minibatch", type=int, default=2, help="Tasks per reflection step" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Seed for selection and minibatches" + ) + parser.add_argument( + "--worker-model", + default=WORKER_MODEL, + help="Model that writes the kernels; the harness's --model is the proposer, " + "as in the paper's proposer/worker split", + ) + parser.add_argument( + "--selection", + choices=["pareto", "best"], + default="pareto", + help="Candidate selection; 'best' mutates the best average instead, which is " + "the naive alternative the paper's 4.3 argues against rather than an ablation " + "it runs", + ) + parser.add_argument( + "--no-side-info", + action="store_true", + help="Score-only feedback: the paper's SI ablation", + ) + parser.add_argument( + "--control-iterations", + type=int, + default=0, + help="Per-task iterations for --single-task; 0 divides --budget across the " + "tasks, which matches the arms on iterations rather than evaluator calls", + ) + parser.add_argument( + "--single-task", + action="store_true", + help="Run the single-task control instead of multi-task search: each task " + "optimized independently at the same per-problem budget", + ) + args = parser.parse_args() + + assert check_reference_agrees() + result = run_kernel(args, random.Random(args.seed)) + report(result, selection=args.selection, side_info=not args.no_side_info) + # No assertion that the score improved: in multi-task mode it cannot go down. The + # headline is a per-task maximum over the pool and the seed is one candidate in it, + # so ``best_score >= seed_score`` holds however badly the search does, and asserting + # it would only look like a check. `report` prints the matched single-artifact + # comparison next to it, which can go down and is the number to read. + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/library.py b/docs/source/llm_examples/optimization/library.py new file mode 100644 index 000000000..9025c5cdc --- /dev/null +++ b/docs/source/llm_examples/optimization/library.py @@ -0,0 +1,711 @@ +"""The optimize_anything engine: reflective Pareto search over a typed artifact. + +Shared machinery for the three runnable examples in this directory, which implement +"optimize_anything: Unified Text Optimization can Outperform Specialized Systems" +(OpenReview 7M28lVzVUq). The paper's observation is that an enormous range of problems +-- a CUDA kernel, a packing algorithm, a scheduling policy, an agent architecture, a +system prompt -- are all *the same problem*: improve an artifact that some evaluator +scores. Its system is GEPA's reflective Pareto search (Agrawal et al., ICLR 2026) +lifted off prompts onto arbitrary artifacts. + +Its Algorithm 1 is short enough to quote, and `optimize_anything` below is it: + + P <- [seed]; evaluate the seed; record per-objective scores + while budget remains: + k <- ParetoSelect(P) # sample in proportion to frontier frequency + M <- a minibatch of 2-3 examples + run candidate k on M, collecting scores *and* side information + k' <- Reflect(k, scores, SI) # the LLM proposes a revision + if k' improves on M: evaluate it fully, admit it, prune dominated candidates + return the best candidate + +Everything except ``Reflect`` is ordinary Python, and that is the point of the split +between this module and its three siblings: what lives here is the whole algorithm, and +what lives in `packing.py`, `prompting.py` and `kernels.py` is only ever an artifact +type, an evaluator, and a prompt. Each of the paper's mechanisms falls out of an +effectful idiom rather than a subsystem: + + * **Side Information is just the evaluator's typed return value.** The paper needs a + ``side_info`` dict and a serializer to carry stack traces, sub-scores and rendered + images to the proposer. Here an evaluator returns an `Evaluation` (score, sub-score + `Metric`s, `Diagnostic`s) and the proposer's prompt splices it with ``{feedback}`` + -- the Encodable bridge already knows how to put a typed value in the model's + context. Because SI is *any* Encodable, a rendered ``PIL.Image`` of the current + packing is SI too, through the same path ``image_input.py`` uses. The paper's SI + ablation is a flag in each script, and it changes one line: which view of the + `Evaluation` the proposer is shown. + + * **The paper's "refiner" is `TenacityRetryer` plus decode-time certification.** It + reports needing a dedicated step for "malformed code blocks, import errors, syntax + issues ... essential for code and agent artifacts where minor formatting errors + cause complete evaluation failure". A code artifact here is a ``Template`` + returning a ``Callable``: the model's source is parsed, type-checked against the + requested signature, compiled, and its own doctests are run at decode time, so a + malformed candidate is fed its own error and revised before it ever reaches the + evaluator. + + * **"Serialize the artifact as a string" is the step you get to skip.** This loop is + generic in the artifact type ``A``. `packing.py` optimizes a callable (in fact a + pair of modules), `prompting.py` and `kernels.py` optimize strings, and nothing in + between ever sees a serialized artifact or needs an adapter. + + * **The three modes are a function signature.** `optimize_anything` takes ``dataset`` + and ``valset``; neither means single-task (the artifact *is* the solution, and the + Pareto objectives are its sub-scores), ``dataset`` alone means multi-task (one + shared frontier, one specialized artifact per task), both means generalization + (search on train, select on held-out val). Pareto selection, minibatching, the + accept-if-improves rule, dominated-candidate pruning and the content-addressed + evaluation cache are plain Python, and stay plain Python. + +The three scripts, each runnable on its own and each documenting what it does and does +not reproduce of its section of the paper: + + * `packing.py` -- circle packing, single-task search over a code artifact (5.3) + * `prompting.py` -- prompt optimization, generalization mode (A.3) + * `kernels.py` -- kernel instructions, multi-task search (5.2) +""" + +import collections.abc +import dataclasses +import inspect +import linecache +import random +import statistics +import typing + +import pydantic.dataclasses + +from effectful.handlers.llm.harness.provision import LiteLLMProvider +from effectful.ops.semantics import handler + +# The one piece of handler boilerplate in this file, and it is load-bearing wherever a +# model sits inside the evaluator. In those domains the paper optimizes an artifact +# *for a specific model* -- a prompt for GPT-4.1-mini, an agent architecture for Gemini +# Flash -- with a strong proposer and a cheap target, the whole point being to lift the +# cheap one. (Its other domains have no target model at all: in circle packing, the +# Optuna comparison and the scheduling algorithms, the artifact is the solution and the +# evaluator is code.) In effectful "run this call on a different model" is a scoped +# handler, so ``worker(...)`` is all the machinery that split needs -- no config +# system, no per-template model registry. +# +# One consequence worth knowing, because it affects how the numbers those domains +# report should be read: this provider does *not* shadow the harness's +# ``TenacityRetryer``. It implements ``completion`` and ``Template.__apply__``, while +# the retryer intercepts ``call_assistant``, which sits between them -- so a worker +# answer that fails to decode is fed its own error and asked again, up to the harness's +# retry limit, exactly as a proposer call would be. Only a failure that exhausts the +# retries reaches the evaluator's ``except`` branch and scores zero. Every score in +# `prompting.py` and `kernels.py` is therefore a score for the artifact *plus that +# repair loop*, not for the artifact alone, and an artifact whose outputs are +# borderline-undecodable is flattered by it. +WORKER_MODEL = "openai/gpt-4.1-mini" + + +def worker(model: str) -> typing.Any: + """Scope a call to the cheap model the artifact is being optimized *for*.""" + return handler(LiteLLMProvider(model=model)) + + +# --------------------------------------------------------------------------- +# Side Information: the evaluator's contract. +# +# The paper's central API claim is that an evaluator returns a score *and* whatever +# diagnostics it can produce, and that the proposer reads them. Here that contract is +# a return type. Note the absence of a dict: values that cross the model boundary use +# ``list``s of small dataclasses, because strict tool schemas reject free-form dicts. +# --------------------------------------------------------------------------- + + +@pydantic.dataclasses.dataclass(frozen=True) +class Metric: + """One sub-score of an evaluation. Higher is always better, by construction -- the + Pareto machinery compares metrics directly, so a "lower is better" quantity is + negated by the evaluator that produces it.""" + + name: str + value: float + + def __str__(self) -> str: + return f"{self.name}={self.value:.6g}" + + +@pydantic.dataclasses.dataclass(frozen=True) +class Diagnostic: + """One piece of Side Information: a named, human-readable explanation of *why* the + artifact scored what it scored -- a violated constraint, a failing test case, a + traceback, a timing. This is the signal the paper argues is the text-optimization + analogue of a gradient.""" + + name: str + detail: str + + def __str__(self) -> str: + return f"{self.name}: {self.detail}" + + +@pydantic.dataclasses.dataclass(frozen=True) +class Evaluation: + """What an evaluator returns: a score, optional sub-scores, and optional Side + Information. ``score_only`` is the paper's SI ablation -- the same evaluation with + its diagnostics withheld, which is all "score-only feedback" means.""" + + score: float + metrics: list[Metric] = dataclasses.field(default_factory=list) + diagnostics: list[Diagnostic] = dataclasses.field(default_factory=list) + + def score_only(self) -> "Evaluation": + return Evaluation(score=self.score) + + def __str__(self) -> str: + lines = [f"score: {self.score:.6g}"] + if self.metrics: + lines.append("metrics: " + ", ".join(str(m) for m in self.metrics)) + lines.extend(f"- {d}" for d in self.diagnostics) + return "\n".join(lines) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Rollout: + """One (example, evaluation) pair handed to the proposer. In single-task mode the + example is the artifact itself, so ``example`` names the task.""" + + example: str + evaluation: Evaluation + + def __str__(self) -> str: + return f"<{self.example}>\n{self.evaluation}\n" + + +# --------------------------------------------------------------------------- +# The generic engine. No LLM appears below this line except through ``proposer``, +# which is the domain's Template call: everything else -- selection, minibatching, +# acceptance, pruning, caching -- is ordinary Python. +# --------------------------------------------------------------------------- + + +class Example(typing.Protocol): + """What the loop needs of a dataset element: a name to key objectives and the + evaluation cache by. Domains supply richer dataclasses; the member is a read-only + property so a frozen dataclass satisfies it.""" + + @property + def name(self) -> str: ... + + +type Evaluator[A, E] = collections.abc.Callable[[A, E | None], Evaluation] +type Reflect[A] = collections.abc.Callable[[A, list[Rollout]], A] + + +def source_of(fn: object) -> str | None: + """The source of a synthesized callable, or ``None`` if it cannot be recovered. + + ``inspect.getsource`` tokenizes the block it finds and raises for more reasons than + its documented OSError/TypeError: a synthesized function whose block ends inside a + multi-line string raises ``tokenize.TokenError``, which is reachable often enough to + end a run. When block extraction fails the whole synthesized module is still sitting + in ``linecache``, and for both of this function's jobs -- keying the cache and + showing the user what the search wrote -- the module is as good as the block. + """ + try: + return inspect.getsource(fn) # type: ignore[arg-type] + except Exception: + pass + try: + path = inspect.getsourcefile(fn) or fn.__code__.co_filename # type: ignore[arg-type, attr-defined] + lines = linecache.getlines(path) + return "".join(lines) or None + except Exception: + return None + + +def artifact_key(artifact: object) -> str: + """Content address of an artifact: its source if it is synthesized code, else its + text. + + Two candidates with the same key have the same evaluation *provided the evaluator + is a function of the artifact and the example and nothing else*. That proviso is + the whole of the cache's soundness, and one domain breaks it: `packing.py` hands + each artifact the best packing found so far, so its score depends on when it ran. + Such a domain must declare a ``state_key`` (see `optimize_anything`), which joins + the key; a content address alone would serve a stale score from before the + incumbent moved. + """ + if callable(artifact): + return source_of(artifact) or repr(artifact) + return str(artifact) + + +@dataclasses.dataclass +class Candidate[A]: + """One artifact in the pool, with its score on every objective. ``refined_on`` + records which examples were in the minibatch that produced it -- the lineage the + multi-task cross-transfer report reads.""" + + index: int + artifact: A + scores: dict[str, float] + parent: int | None + generation: int + refined_on: list[str] = dataclasses.field(default_factory=list) + + @property + def average(self) -> float: + return statistics.fmean(self.scores.values()) if self.scores else 0.0 + + def dominates(self, other: "Candidate[A]", objectives: list[str]) -> bool: + """Pareto dominance: at least as good everywhere, strictly better somewhere.""" + at_least = all(self.scores[j] >= other.scores[j] for j in objectives) + strictly = any(self.scores[j] > other.scores[j] for j in objectives) + return at_least and strictly + + +def pareto_frontier[A]( + pool: list[Candidate[A]], objectives: list[str] +) -> list[Candidate[A]]: + """The non-dominated candidates: everything that is the best at *something*.""" + return [ + c + for c in pool + if not any(o.dominates(c, objectives) for o in pool if o is not c) + ] + + +def pareto_select[A]( + pool: list[Candidate[A]], objectives: list[str], rng: random.Random +) -> Candidate[A]: + """GEPA's selection rule, which the paper adopts verbatim: among non-dominated + candidates, sample in proportion to how many objectives a candidate is *best* at. + A specialist that wins one objective stays reachable; a generalist that wins many + is reached often.""" + frontier = pareto_frontier(pool, objectives) + weights = [ + sum( + 1 + for j in objectives + if c.scores[j] >= max(f.scores[j] for f in frontier) - 1e-12 + ) + for c in frontier + ] + if not any(weights): # degenerate scores -- fall back to uniform + return rng.choice(frontier) + return rng.choices(frontier, weights=weights, k=1)[0] + + +def best_select[A]( + pool: list[Candidate[A]], objectives: list[str], rng: random.Random +) -> Candidate[A]: + """Always mutate the best average, collapsing the frontier's complementary + strengths into one number. + + This is the naive alternative the paper's 4.3 argues against -- "averaging hides + which aspects are strong and which are weak" -- not an ablation it runs. The + scripts expose it as ``--selection best`` so the argument can be checked here, and + none of them has yet spent the budget to check it. + """ + return max(pool, key=lambda c: c.average) + + +@dataclasses.dataclass +class Step: + """One iteration of the loop, kept for the printed trace.""" + + iteration: int + parent: int + before: float + after: float + accepted: bool + best: float + + +@dataclasses.dataclass +class Result[A, E]: + """The outcome of a run: the pool (so the surviving frontier can be inspected), + the trace, and the mode-appropriate answer.""" + + mode: str + pool: list[Candidate[A]] + history: list[Step] + objectives: list[str] + seed_score: float + best: Candidate[A] + best_score: float + per_task: list[tuple[str, Candidate[A], float]] = dataclasses.field( + default_factory=list + ) + evaluations: int = 0 + proposals: int = 0 + + def frontier(self) -> list[Candidate[A]]: + return pareto_frontier(self.pool, self.objectives) + + +def mode_of(dataset: object, valset: object) -> str: + if dataset is None: + return "single-task" + return "generalization" if valset is not None else "multi-task" + + +def optimize_anything[A, E: Example]( + *, + evaluator: Evaluator[A, E], + proposer: Reflect[A], + seed: A | None = None, + bootstrap: collections.abc.Callable[[str], A] | None = None, + objective: str | None = None, + dataset: list[E] | None = None, + valset: list[E] | None = None, + budget: int = 8, + minibatch_size: int = 2, + selection: str = "pareto", + use_side_info: bool = True, + rng: random.Random | None = None, + task_name: str = "task", + state_key: collections.abc.Callable[[], str] | None = None, +) -> Result[A, E]: + """The paper's Algorithm 1, generic in the artifact type. + + The mode is a function of the arguments and nothing else: no ``dataset`` is + single-task (objectives are the artifact's sub-score metrics), ``dataset`` alone is + multi-task (objectives are the tasks; each task selects its own artifact off the + shared frontier), and ``dataset`` + ``valset`` is generalization (search on train, + select on held-out val). Seedless mode -- ``seed=None`` with an ``objective`` and a + ``bootstrap`` -- lets the model write candidate zero. + + ``state_key`` names whatever *else* an evaluation depends on. It exists because the + paper hands each circle-packing artifact the best solution found so far + (``main(timeout, current_best_solution)``), which means the evaluator stops being a + pure function of the artifact and a content-addressed cache silently starts serving + stale scores. A domain that threads state like that says so here; every other + domain leaves it ``None`` and keeps the plain content address. + """ + rng = rng or random.Random(0) + mode = mode_of(dataset, valset) + cache: dict[tuple[str, str, str], Evaluation] = {} + counters = {"evaluations": 0, "proposals": 0} + + def evaluate(artifact: A, example: E | None) -> Evaluation: + """Content-addressed evaluation: the paper's caching, in three lines. It earns + its keep because an evaluation can itself be an LLM call.""" + key = ( + artifact_key(artifact), + example.name if example else task_name, + state_key() if state_key else "", + ) + if key not in cache: + counters["evaluations"] += 1 + cache[key] = evaluator(artifact, example) + return cache[key] + + def score_on(artifact: A, examples: list[E] | list[None]) -> dict[str, float]: + """Per-objective scores. With a dataset the objectives are the examples; with + none, they are the single evaluation's sub-score metrics (the paper: "single- + task search admits only one data point, so per-example tracking reduces to + per-metric tracking").""" + if dataset is None: + evaluation = evaluate(artifact, None) + scores = {m.name: m.value for m in evaluation.metrics} + # The headline is only its own objective when the evaluator offers nothing + # finer; adding it alongside metrics that already contain it would just + # double that dimension's weight in the frontier-frequency sampling. + if not scores: + scores["score"] = evaluation.score + return scores + return { + e.name: evaluate(artifact, e).score for e in typing.cast(list[E], examples) + } + + def headline(candidate: Candidate[A]) -> float: + """The number a human reads. Averaging heterogeneous sub-scores is meaningless + in single-task mode, where the artifact's own score is the answer; with a + dataset the objectives *are* the per-example scores, so the average is right.""" + if dataset is not None: + return candidate.average + return candidate.scores.get("max_score") or candidate.scores.get("score", 0.0) + + pool_examples: list[E] | list[None] = ( + typing.cast(list[E] | list[None], dataset) if dataset is not None else [None] + ) + + # --- candidate zero ----------------------------------------------------- + if seed is None: + if bootstrap is None or objective is None: + raise ValueError( + "seedless mode needs both an ``objective`` and a ``bootstrap``" + ) + print("[seedless] bootstrapping candidate 0 from the objective ...") + seed = bootstrap(objective) + root = Candidate( + index=0, + artifact=seed, + scores=score_on(seed, pool_examples), + parent=None, + generation=0, + ) + pool: list[Candidate[A]] = [root] + # Candidate indices come from a monotonic counter, not ``len(pool)``: pruning + # removes dominated candidates, so a length-derived index would be reused and the + # trace's parent links would silently point at the wrong artifact. + minted = 1 + objectives = sorted(root.scores) + history: list[Step] = [] + best_so_far = headline(root) + print( + f"[{mode}] {len(objectives)} objective(s): {', '.join(objectives)}\n" + f"[seed] score {best_so_far:.8g}" + ) + + select = pareto_select if selection == "pareto" else best_select + + # --- the loop ----------------------------------------------------------- + for iteration in range(1, budget + 1): + parent = select(pool, objectives, rng) + + # A minibatch of 2-3 examples, not the whole set: the paper's second Pareto + # ingredient, so reflection is focused instead of trying to fix everything. + minibatch: list[E] | list[None] + if dataset is None: + minibatch = [None] + else: + minibatch = rng.sample(dataset, k=min(minibatch_size, len(dataset))) + + rollouts = [ + Rollout( + example=e.name if e is not None else task_name, + evaluation=( + evaluate(parent.artifact, e) + if use_side_info + else evaluate(parent.artifact, e).score_only() + ), + ) + for e in minibatch + ] + before = statistics.fmean(r.evaluation.score for r in rollouts) + + # The only LLM call in the loop: reflect over the minibatch and its SI. + counters["proposals"] += 1 + try: + child_artifact = proposer(parent.artifact, rollouts) + except Exception as exc: + # Retries are exhausted, so the decode-time gate has rejected this proposal + # for good and the iteration is lost. The reason is worth printing: it is + # usually the artifact's own doctests failing, which is the paper's refiner + # step doing its job where you can see it. + reason = " ".join(str(exc).split())[:200] + print( + f" iter {iteration}: proposal rejected at decode " + f"({type(exc).__name__}: {reason})" + ) + history.append( + Step(iteration, parent.index, before, before, False, best_so_far) + ) + continue + + after = statistics.fmean(evaluate(child_artifact, e).score for e in minibatch) + + accepted = after > before + if accepted: + # Only now pay for a full evaluation -- the paper's ordering. + child = Candidate( + index=minted, + artifact=child_artifact, + scores=score_on(child_artifact, pool_examples), + parent=parent.index, + generation=parent.generation + 1, + refined_on=[r.example for r in rollouts], + ) + minted += 1 + pool.append(child) + kept = pareto_frontier(pool, objectives) + dropped = len(pool) - len(kept) + pool = kept + best_so_far = max(best_so_far, max(headline(c) for c in pool)) + else: + dropped = 0 + + history.append( + Step(iteration, parent.index, before, after, accepted, best_so_far) + ) + print( + f" iter {iteration}: parent #{parent.index} (gen {parent.generation}) " + f"minibatch {before:.8g} -> {after:.8g} " + f"{'ACCEPTED' if accepted else 'rejected'}" + + (f", pruned {dropped}" if dropped else "") + + f", best {best_so_far:.8g}" + ) + + # --- what "best" means depends on the mode ------------------------------ + per_task: list[tuple[str, Candidate[A], float]] = [] + if mode == "generalization": + # Search used the train set; the answer is whatever generalizes. Only the + # frontier is re-evaluated on val, since evaluation costs model calls. + frontier = pareto_frontier(pool, objectives) + val_scores = { + c.index: statistics.fmean( + evaluate(c.artifact, e).score for e in typing.cast(list[E], valset) + ) + for c in frontier + } + best = max(frontier, key=lambda c: val_scores[c.index]) + best_score = val_scores[best.index] + seed_score = statistics.fmean( + evaluate(root.artifact, e).score for e in typing.cast(list[E], valset) + ) + elif mode == "multi-task": + # N specialized artifacts: each task picks its own best off the shared + # frontier, which is exactly where cross-transfer shows up. The run's score is + # therefore the mean over those per-task winners, not any single candidate's + # average -- the paper's "each task independently selects its own best + # candidate from the frontier". Scoring one artifact across all tasks would + # understate multi-task mode by construction, since specializing is the point. + for e in typing.cast(list[E], dataset): + winner = max(pool, key=lambda c: c.scores[e.name]) + per_task.append((e.name, winner, winner.scores[e.name])) + best = max(pool, key=lambda c: c.average) + best_score = statistics.fmean(score for _, _, score in per_task) + seed_score = root.average + else: + best = max(pool, key=headline) + best_score = headline(best) + seed_score = headline(root) + + return Result( + mode=mode, + pool=pool, + history=history, + objectives=objectives, + seed_score=seed_score, + best=best, + best_score=best_score, + per_task=per_task, + evaluations=counters["evaluations"], + proposals=counters["proposals"], + ) + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- + + +def transfer_note(result: "Result") -> str: + """How many per-task winners were last refined while the proposer was looking at a + different task -- stated against the count chance alone would produce. + + Two things make the raw count meaningless on its own. A winner counts as + "transferred" unless its own task was in the minibatch that produced it, so under a + null where any candidate is equally likely to win any task the expected count is + already ``tasks * (1 - minibatch/tasks)`` -- 3.0 on five tasks with a two-task + minibatch, which is most of the range the statistic can take. And a task whose + winner is the *seed* was never refined at all: it has no minibatch, so it cannot + have transferred, and counting it as one turns a run with no transfer whatsoever + into "4/5, above chance". Both the count and the null are therefore taken over the + refined winners only, each contributing its own minibatch size rather than a pooled + mean, and seed winners are reported separately. + + It remains one run, and it looks only at the last refinement rather than at full + ancestry. + """ + winners = result.per_task + seeds = [name for name, candidate, _ in winners if not candidate.refined_on] + refined = [(name, c) for name, c, _ in winners if c.refined_on] + transferred = [name for name, c in refined if name not in c.refined_on] + # Per candidate, since minibatches can differ in size: the chance of *not* drawing + # this task into the minibatch that produced its winner. + expected = sum(1.0 - len(c.refined_on) / len(winners) for _, c in refined) + if not refined: + return ( + f"\nCross-task transfer: not measurable -- all {len(winners)} winners are " + f"the seed, which was never refined on any task" + ) + verdict = ( + "above chance" + if len(transferred) > expected + else "at or below chance, so this run is no evidence of transfer" + ) + return ( + f"\nCross-task transfer: {len(transferred)}/{len(refined)} *refined* winners " + f"were last refined while looking at a different task" + + (f" ({', '.join(transferred)})" if transferred else "") + + f"; chance alone would give {expected:.1f} -- {verdict}" + + ( + f". The other {len(seeds)} winner(s) ({', '.join(seeds)}) are the " + f"unrefined seed and are excluded from both figures" + if seeds + else "" + ) + ) + + +def report( + result: "Result", + *, + selection: str, + side_info: bool, + notes: collections.abc.Sequence[str] = (), + render_artifact: collections.abc.Callable[[typing.Any], str] | None = None, +) -> None: + """Print what a run did: the trace's summary, the surviving frontier, per-task + winners where the mode has them, and the winning artifact. + + ``notes`` and ``render_artifact`` are how a domain adds its own reading without the + library having to know about it -- `packing.py` uses them to count which of its two + modules the frontier came from, and to print a two-part artifact. + """ + print("\n" + "=" * 72) + print( + f"mode: {result.mode} | selection: {selection} | " + f"side information: {'on' if side_info else 'off'}" + ) + print( + f"seed {result.seed_score:.6g} -> best {result.best_score:.6g} " + f"after {len(result.history)} iterations " + f"({result.proposals} proposals, {result.evaluations} evaluations)" + ) + accepted = sum(step.accepted for step in result.history) + print(f"accepted proposals: {accepted}/{len(result.history)}") + # The multi-task headline is a per-task maximum over the pool while the seed is a + # single candidate, so it cannot be negative and is biased upward by the size of the + # pool. The best single artifact's mean is the matched comparison: one artifact + # against one artifact, on the same tasks. It is only meaningful when every + # candidate was scored on the same objectives, which is false for an arm that + # aggregates independent runs. + comparable = result.pool and all( + set(c.scores) == set(result.objectives) for c in result.pool + ) + if result.per_task and comparable: + best_single = max(c.average for c in result.pool) + print( + f"best single artifact (matched comparison, since the headline above is a " + f"per-task maximum over {len(result.pool)} candidates): {best_single:.6g}" + ) + + if result.pool: + frontier = result.frontier() + print(f"\nPareto frontier ({len(frontier)} candidate(s) survive):") + for c in sorted(frontier, key=lambda c: -c.average): + scores = ", ".join(f"{k}={v:.4g}" for k, v in sorted(c.scores.items())) + print(f" #{c.index} (gen {c.generation}, parent {c.parent}): {scores}") + else: + print("\n(no shared frontier: this run is several independent searches)") + + if result.per_task: + print("\nPer-task winners:") + for name, candidate, score in result.per_task: + print(f" {name}: candidate #{candidate.index} scored {score:.4g}") + # Cross-task transfer is only a question where there was one frontier to + # transfer across; for independent runs it is zero by construction, and saying + # so as though it were a measurement would be worse than not saying it. + if result.pool: + print(transfer_note(result)) + + for note in notes: + print(f"\n{note}") + + print("\nBest artifact:") + artifact = result.best.artifact + if render_artifact is not None: + print(render_artifact(artifact)) + elif callable(artifact): + print((source_of(artifact) or repr(artifact)).rstrip()) + else: + print(artifact) diff --git a/docs/source/llm_examples/optimization/packing.py b/docs/source/llm_examples/optimization/packing.py new file mode 100644 index 000000000..165f54a28 --- /dev/null +++ b/docs/source/llm_examples/optimization/packing.py @@ -0,0 +1,1376 @@ +"""Circle packing: single-task search over a code artifact (optimize_anything 5.3). + +Pack ``n`` non-overlapping circles into the unit square so the sum of their radii is as +large as possible. The artifact *is* the solution here -- there is no dataset, and the +evaluator scores the candidate directly -- which is the paper's single-task mode, and +the mode AlphaEvolve and OpenEvolve operate in. + +This script is set up to be a genuine attempt at the paper's result rather than an +illustration of the loop, so it follows 5.3 and Appendix G closely: + + * The artifact has the paper's signature, ``pack(n, time_budget, current_best)`` + (its evolved packer is ``main(timeout, current_best_solution)``, Appendix K.6), and + is handed the best packing found so far to polish. That is what lets a search of a + few dozen evaluations reach a competitive number instead of restarting each time -- + and it is why `optimize_anything` needs a ``state_key``, since an artifact's score + now depends on when it ran. + * It may use whatever numeric libraries are actually installed, because the paper's + winner is an LP over radii whose dual variables give gradients for a local + optimizer over centres -- unreachable in the standard library. The prompt reports + what ``importlib`` finds rather than a fixed list, so the example still runs where + scipy is absent. It reports *only* that; see `numeric_toolbox`. + * The Pareto objectives are Mechanism 3's run-distribution metrics rather than a + single number, which is what keeps structurally different packers alive on the + frontier -- three of them here rather than the paper's four, for the reason + `trajectory_metrics` gives. + * The search evolves *two* modules on one shared frontier, the packer and a refiner + instruction, which is Mechanism 2's leapfrogging. This needs nothing from the + engine: a candidate is a `PackSystem`, and "which module to mutate" is a branch in + this script's proposer. + +The evaluator is deterministic Python throughout, so no model sits in the scoring path +and every number in the trace is measured. + +Demonstrates: +- A ``Template`` returning a ``Callable`` whose *own* doctests are the decode-time + contract: a packer that does not return ``n`` feasible circles is fed its error by + ``TenacityRetryer`` and never reaches the evaluator -- the paper's refiner stage, + for free +- Side Information as a typed value: geometric diagnostics, sub-scores, the spread + across repeated runs, and with ``--visual-si`` a rendered ``PIL.Image`` of the + current packing, all spliced into the proposer's prompt through the Encodable bridge +- Multi-module search with no engine support at all -- code and refiner instruction + compete on the one frontier +- Search over state the evaluator depends on, declared through ``state_key`` +- Scoring that cannot be gamed by a tolerance: the sum of radii is measured *after* + shrinking the packing to exact feasibility +- A control that says what the search is worth: ``--baseline`` runs the same problem + for the same wall-clock with no model anywhere in it + +What the numbers are +-------------------- + +Measured on 2026-07-30 with gpt-5.5 proposing, on the paper's own instance +(``--num-circles 26 --time-budget 20 --budget 10``), one run per configuration, each +from the 6x6 grid seed at 2.1666667: + + * default -- Pareto selection, side information on: **2.6359831**, 3 of 10 proposals + accepted, 15 evaluations + * ``--no-side-info``: 2.6319369, 2 of 10 accepted, 19 evaluations + * ``--selection best``: 2.6317302, 5 of 10 accepted, 17 evaluations + +The paper reports 2.63598 on this instance, against 2.635 for AlphaEvolve and 2.6307 +for OpenEvolve at 200 evaluations, so the default arm matches its value to every digit +it gives. A +repeat of that arm reached the same 2.6359831 after four iterations before the +wall-clock hazard below wedged it at iteration 5; two runs landing on exactly that value +from different proposals is what a real local optimum looks like, and it is the +strongest evidence here that the artifact solves the problem rather than reciting a +published answer for 26. Four things have to be said before any of this is read as a +reproduction. + +*Most of the distance is scipy's, not the search's.* ``--baseline`` runs warm-started +random-restart SLSQP with no LLM in the loop for the same packer wall-clock the search +spends (10 iterations x 3 repeats x 20s = 600s), and it reaches **2.6342924** in 13116 +restarts -- past OpenEvolve at 200 evaluations, within 0.0007 of AlphaEvolve, and above +both ablation arms. The default arm's first accepted proposal scores +2.6342924 exactly: the first competent artifact the search writes is doing what the +control does, digit for digit, and the whole remaining margin of the run is 0.0017. So +the honest statement of this domain's result is that a reflective search whose artifacts +call a constrained optimizer beat a plain call to the same optimizer by about 0.0017 at +matched wall-clock. The paper reports no such control, which is why its own margin over +specialized systems is not attributable either. + +*The run's number and the artifact's are different numbers, and the gap varies by arm.* +Every candidate is handed the best packing found so far and told never to return worse, +so a score accumulates the work of everything before it -- a candidate returning its +input unchanged is recorded at the full incumbent value. This is the paper's setup, not +a deviation from it: its evolved packer takes ``current_best_solution`` too, so its +2.63598 is a trajectory number in the same way. `cold_start_note` re-runs the winner +with no incumbent and prints both. The default arm's winner scores 2.6359831 cold +against the run's 2.6359831 -- it inherited nothing and reaches the headline from the +grid on its own. ``--selection best``'s winner scores 2.6319369 cold against a run +number of 2.6317302, marginally *better* alone than in the run. ``--no-side-info``'s +winner scores 2.5416318 cold against 2.6319369, so 0.09 of its score is other +candidates' work rather than its own, and that arm's headline is the least attributable +of the three. + +*Score-only feedback costs almost nothing here, and the paper's ablation figure does not +reproduce.* ``--no-side-info`` reaches 2.6319369, which is 99.85% of the side-information +arm (99.1% of the distance from the seed), against the 93.96% the paper's Table 4 +reports. The direction is the paper's and the magnitude is not, and the trace says why: +the first accepted proposal in every arm jumps from the grid to a warm-started SLSQP +restart loop and lands within 0.005 of the best number any arm reaches, after which all +of them grind in the fourth decimal. Diagnostics naming which circles are jammed cannot +be worth much when the remaining headroom is 0.004 and the artifact's own optimizer is +already searching it. That is a fact about this domain rather than a refutation of the +paper's: on a task whose ceiling is one competent artifact away from the seed, the SI +ablation has almost nothing to measure. + +*Greedy selection is not distinguishable from Pareto here, because the frontier never +holds more than one candidate.* ``--selection best`` reached 2.6317302 against the +default arm's 2.6359831. That looks like support for 4.3's argument against collapsing +the frontier to an average, and it is not: every one of these runs ends "Pareto frontier +(1 candidate(s) survive)", and every accepted proposal prunes exactly one candidate. The +three objectives (max, mean and worst over the repeats) move together for packers this +close to deterministic, so dominance is total, and Pareto selection spends the run +choosing from a pool of one. With a single run per arm and no variance estimate, a 0.004 +difference between two configurations that both reduce to "mutate the only candidate +there is" measures nothing about the selection rule, and the mechanism the difference +would have to come from -- structurally different packers kept alive by complementary +strengths -- never appears in the trace. 4.3 is untested here, not confirmed. + +`module_note` gives Mechanism 2 the same treatment. Across the three runs the refiner +module's accepted proposals carry mean minibatch gains of +0.4676, +0.2313 and +0.2320; +the code module's carry +0.000845 and +0.000353, and in the score-only arm it had +nothing accepted at all. The refiner wins the one move that matters, off the grid seed, +and the code module grinds out everything after it in the fourth decimal and beyond, +which makes the two modules' gain figures a statement about when each ran rather than +about how good either is. The modules +do alternate -- accepted gains arrive as ``refiner -> code -> code`` in the default arm +and ``refiner -> code -> code -> refiner -> code`` in ``--selection best``, three +handovers -- but a handover means only that the other module produced the next accepted +gain, not that it was ahead of its partner. What the counts do establish is that the +second module is not decoration: it produced the first accepted gain in all three runs +and the winning artifact in the score-only arm. + +Iterations 7 and 9 of the ``--selection best`` run both read ``2.6317302 -> 2.6317302 +ACCEPTED``, as does iteration 3 of the default arm at 2.6342924. The accept gate is a +bare ``after > before``, so a child that clamped to the incumbent and improved it in the +eighth decimal is admitted and its parent pruned. That is Algorithm 1's accept rule as +written, and it is the mechanism by which a candidate contributing nothing carries the +run's whole score forward. + +The winning artifact is an algorithm rather than a remembered answer, with one +qualification worth stating. It builds the 4n wall constraints and n(n-1)/2 separation +constraints programmatically over 3n variables, hands them to SLSQP warm-started from +the incumbent packing, repairs every iterate to exact feasibility before scoring it, and +keeps the best; it is recognisably the same program as `baseline_packing`, which is the +point above, and it contains no coordinate table. It does contain +``random.Random(15 if n == 26 else 1000003 + 7919 * n)`` -- a restart seed picked for the +instance it was asked about. So its zero cold-start gap says it reliably reproduces its +own lucky restart sequence at n=26, which is a weaker claim than reliably finding +2.6359831. Run directly at a size nobody asked it about, and given the same 20s the +search gave it, it scores 2.7827752 for n=29 against 2.4166667 for the grid, so it is an +algorithm on the evidence rather than on its author's word -- but see the note on +`generality_check` below, which does not give it that budget and concludes the opposite. +""" + +# Simplifications vs. the source: +# - Budget is counted in optimizer iterations rather than metric calls or dollars; the +# paper spends 63 evaluations and $3.18 on this domain. +# - Three Mechanism-3 objectives, not the paper's four: it names them without defining +# them, and two of the four readings attempted here measured noise or rewarded doing +# nothing (see `trajectory_metrics`). +# - The two modules alternate by a coin flip (``--code-share``); the paper does not say +# how it splits attention between them. There is no per-module score to plot, so +# Mechanism 2's leapfrogging curve is not reproduced -- only which module produced +# each accepted gain (see `PackSystem` and `module_note`). +# - `generality_check` gives the packer 0.5s, and an artifact that honours a short budget +# by returning its safe fallback is reported as "a table rather than an algorithm" on +# that basis. It says exactly that about the winning artifact above, which scores +# 2.7827752 at n=29 against the grid's 2.4166667 when given the run's own 20s. So this +# diagnostic currently distinguishes "bails out when rushed" from "hardcodes an +# answer" not at all, and the proposer is shown the wrong conclusion every iteration. +# Raising its budget would cost a full extra packer run per evaluation, which is why +# the cheap version is here; the trade is not free either way. +# - One run per configuration and no variance estimate, on a domain where the three +# configurations measured span 2.6317 to 2.6360 -- a range as wide as the whole gap +# between the published systems the headline is compared against. +# - The wall-clock backstop in `_run_packer` is best-effort, not a guarantee. It is a +# Python-level ``SIGALRM``, and the handler only runs when the interpreter next gets +# control, so a synthesized packer sitting in a long C call or one that has moved work +# into a subprocess can hang a run indefinitely -- which does happen. A real bound +# needs process isolation, which this example does not do. +# - The accept gate is ``after > before`` with no minimum improvement, so a proposal +# worth a billionth of the score is accepted and prunes its parent. Adding a threshold +# would depart from Algorithm 1 as written, so it is documented rather than changed. +# - No island model / MAP-Elites, which the paper also drops. + +import argparse +import collections +import collections.abc +import dataclasses +import importlib.util +import io +import math +import random +import signal +import statistics +import threading +import time +import traceback +import typing + +import pydantic.dataclasses +from PIL import Image + +from docs.source.llm_examples.optimization.library import ( + Diagnostic, + Evaluation, + Metric, + Result, + Rollout, + artifact_key, + optimize_anything, + report, + source_of, +) +from effectful.handlers.llm import Agent, Template + + +@pydantic.dataclasses.dataclass(frozen=True) +class Circle: + """A circle in the unit square: centre and radius.""" + + x: float + y: float + r: float + + def __str__(self) -> str: + return f"({self.x:.4f}, {self.y:.4f}) r={self.r:.4f}" + + +# ``pack(n, time_budget, current_best)`` is the paper's artifact signature -- its +# evolved circle packer is ``main(timeout, current_best_solution)`` (Appendix K.6). The +# artifact is handed its own time budget *and* the best packing found so far, so a +# candidate can polish the incumbent instead of starting over every time. That is what +# lets a search of a few dozen evaluations reach a competitive number, and it is why +# ``optimize_anything`` needs a ``state_key``: with the incumbent threaded through, an +# artifact's score depends on when it ran. +type Packer = collections.abc.Callable[[int, float, list[Circle] | None], list[Circle]] + + +def feasible(circles: collections.abc.Sequence[Circle], tol: float = 1e-9) -> bool: + """True when every circle lies inside the unit square and no two overlap. + + In the synthesized packer's lexical scope, so the doctests the model must write + can call it -- that is what makes "the artifact obeys its contract" checkable at + decode time rather than at evaluation time. + + >>> feasible([Circle(0.25, 0.25, 0.25), Circle(0.75, 0.75, 0.25)]) + True + >>> feasible([Circle(0.5, 0.5, 0.6)]) + False + """ + return worst_violation(circles) <= tol + + +def worst_violation(circles: collections.abc.Sequence[Circle]) -> float: + """How badly the packing breaks its constraints: the largest overlap depth or + out-of-square excursion, and 0.0 for a feasible packing. + + >>> worst_violation([Circle(0.5, 0.5, 0.5)]) + 0.0 + >>> round(worst_violation([Circle(0.5, 0.5, 0.5), Circle(0.5, 0.5, 0.5)]), 6) + 1.0 + """ + worst = 0.0 + for c in circles: + if c.r <= 0.0: + worst = max(worst, 1.0 - c.r) + worst = max(worst, c.r - c.x, c.r - c.y, c.x + c.r - 1.0, c.y + c.r - 1.0) + for i, a in enumerate(circles): + for b in circles[i + 1 :]: + worst = max(worst, a.r + b.r - math.hypot(a.x - b.x, a.y - b.y)) + return max(0.0, worst) + + +def total_radius(circles: collections.abc.Sequence[Circle]) -> float: + """The reported sum of the radii -- what the packer claims. + + >>> total_radius([Circle(0.25, 0.25, 0.25), Circle(0.75, 0.75, 0.25)]) + 0.5 + """ + return sum(c.r for c in circles) + + +# The paper's winning packer is a bilevel optimizer: an LP over radii whose duals give +# exact gradients for L-BFGS-B over centres, plus CMA-ES exploration (Appendix K.6). +# None of that is reachable in the standard library, so the proposer is told what is +# actually importable here rather than a fixed list -- scipy and numpy arrive in this +# repo transitively, and the example still runs where they do not. +NUMERIC_LIBRARIES = [ + name + for name in ("numpy", "scipy.optimize", "scipy.spatial") + if importlib.util.find_spec(name) is not None +] + + +def numeric_toolbox() -> str: + """What the synthesized packer may import, as a sentence for the prompt. + + Only that, and the restraint is the point. The paper's *result* on this domain is a + particular algorithm -- an exact LP over the radii for fixed centres, its duals as + gradients on the centres, the two alternated (Appendix K.6) -- and naming it in the + prompt that asks the search to find it turns a search into a transcription. What + belongs here is the part the proposer cannot discover for itself, because it cannot + import anything to find out: which libraries exist. + """ + if not NUMERIC_LIBRARIES: + return ( + "Only the Python standard library is available -- no numpy, no scipy -- so " + "write the numerics yourself." + ) + return ( + f"These numeric libraries are installed and you may import them: " + f"{', '.join(NUMERIC_LIBRARIES)}." + ) + + +def feasible_scale(circles: collections.abc.Sequence[Circle]) -> float: + """The largest factor ``s <= 1`` for which scaling every radius by ``s`` makes the + packing exactly feasible -- 1.0 for a packing with room to spare, 0.0 for one that + cannot be rescued. + + The score is ``s * total_radius``, and that is deliberate. Scoring the *reported* + radii against a tolerance invites the artifact to overshoot by just under it: a + packer that adds 4e-10 to every radius sits inside a 1e-9 feasibility check and + collects the difference, and a search will find that before it finds a better + packing. Shrinking to exact feasibility instead of thresholding removes the + incentive -- an inflated radius is scaled straight back out, and it drags every + other circle down with it -- and it replaces the feasible/infeasible cliff with a + gradient the proposer can actually climb. + + >>> feasible_scale([Circle(0.25, 0.25, 0.25), Circle(0.75, 0.75, 0.25)]) + 1.0 + >>> feasible_scale([Circle(0.5, 0.5, 1.0)]) + 0.5 + """ + scale = 1.0 + for c in circles: + wall = min(c.x, c.y, 1.0 - c.x, 1.0 - c.y) + if c.r <= 0.0 or wall <= 0.0: + return 0.0 # a non-positive radius or a centre outside the square + scale = min(scale, wall / c.r) + for i, a in enumerate(circles): + for b in circles[i + 1 :]: + scale = min(scale, math.hypot(a.x - b.x, a.y - b.y) / (a.r + b.r)) + return max(0.0, min(1.0, scale)) + + +@dataclasses.dataclass(frozen=True) +class PackSystem: + """The artifact of the packing domain: *two* modules, not one. + + The paper optimizes the code artifact and a refiner prompt together on a single + shared Pareto front (its Mechanism 2, "multi-module Pareto leapfrogging"), and + credits that coordination for the circle-packing result: the refiner discovers an + LP-based approach while the code module is still a weak heuristic, the code module + absorbs it and catches up, the refiner pushes further with sequential LP, the code + absorbs that too. Each module's advance is the foundation for the other's next one. + + Modelling that needs nothing from the engine: a candidate is a ``PackSystem``, and + the domain's proposer decides on each iteration which module to mutate -- rewrite + the packer directly, or rewrite the refiner instruction and apply it to the packer. + Both paths produce a new ``PackSystem`` that lands on the same frontier. ``origin`` + records which module produced it, so the leapfrogging is countable afterwards + (`module_note`); it is deliberately left out of ``__str__`` so it does not perturb + the cache key. + + Two things the paper credits to this mechanism are *not* reachable here: + + * The paper's "a failed code mutation is recovered rather than lost, because the + refiner can rewrite it" cannot happen in this implementation. The refiner is + applied to a parent drawn from the pool, and the accept gate never admits a + candidate that scored zero, so the packer handed to the refiner has always + already passed. Only a broken *seed* could be repaired this way. + * The leapfrogging the paper measures is a per-module score curve -- code at 0.98 + while the refiner is at 1.93, then the reverse. There is no per-module score to + plot here: the refiner only ever affects the world through the packer it + rewrites, so a ``PackSystem`` has one score and the two modules share it. What + `module_note` can honestly report is which module produced each accepted gain + and whether the two alternate, which is the observable shadow of leapfrogging + rather than the measurement itself. + """ + + packer: Packer + refiner: str + origin: str = "seed" + + def __str__(self) -> str: + return ( + f"\n{artifact_key(self.packer)}\n\n" + f"\n{self.refiner}\n" + ) + + +# The refiner module's starting point: a generic repair instruction, which the search +# is free to turn into something specific about packing. +SEED_REFINER = ( + "Look at the diagnostics, find the single change that would raise the score the " + "most, and make it." +) + + +class Proposer(Agent): + """You are a reflective optimizer. You are shown the current artifact, the score + it achieved, and diagnostic side information explaining *why* it scored that way, + and you return a strictly better artifact. You do not mutate blindly: you first + read the diagnostics to decide which failure mode is costing the most, then you + make the change that addresses it -- and you are willing to replace the whole + approach with a different one when the diagnostics say the current approach has + saturated.""" + + @Template.define + def propose_packer( + self, current: Packer, feedback: list[Rollout], n: int, toolbox: str + ) -> Packer: + """Write an improved ``pack(n, time_budget, current_best)`` that packs {n} + non-overlapping circles into the unit square [0,1]x[0,1], maximizing the SUM OF + THE RADII. The circles may have different radii. Return a list of + ``Circle(x, y, r)``. + + The current artifact scored as follows. Read the diagnostics before you write + anything: they tell you which circles are jammed, where the slack is, how the + score varied across repeated runs, and whether the packing is even feasible. + + + {feedback} + + + Your arguments: + - ``n``: how many circles. Read it; never hardcode a table for one size. + - ``time_budget``: seconds you may spend. Poll ``time.monotonic()`` and return + your best packing before it expires. Spend it -- returning early wastes + search you were given. + - ``current_best``: the best packing found so far (a list of ``Circle``), or + ``None`` on the first call. POLISH IT. Starting from the incumbent and + improving it is how a handful of evaluations reaches a strong number; + restarting from scratch every time throws that away. Keep it as one of your + starting configurations even when you also try fresh ones, and never return + something worse than what you were handed. + + {toolbox} + + Other constraints: + - You are scored on the sum of radii AFTER the whole packing is shrunk to exact + feasibility, so an overlap costs you proportionally and padding a radius to + sit just inside a tolerance gains you nothing: it is scaled straight back + out, and it shrinks every other circle with it. + - The evaluator runs you several times and looks at the distribution, so + randomness is fine and a stable, repeatable method is worth more than a lucky + one. + - It must work for ANY ``n``: the evaluator also reports how you do on a size + you were not asked about. + + Your function's docstring MUST contain doctests certifying the contract, and + they are run before your artifact is accepted -- this is the decode-time gate + the paper builds a separate "refiner" stage for. ``Circle``, ``feasible``, + ``total_radius`` and ``worst_violation`` are in scope. Write at least a + doctest that binds ``cs = pack(5, 0.5, None)`` and checks + ``len(cs) == 5 and feasible(cs)``, prefixing each input line with the doctest + prompt (three ``>`` characters and a space; it is spelled out rather than + shown so that this instruction is not itself collected as a test). + + The doctest below certifies the same contract on the other decode path, where + the harness synthesizes this Template's body: it calls this Template + recursively -- routed to your own submission, so it costs nothing -- and runs + the packer that comes back. + + >>> _packer = Proposer().propose_packer(seed_packer, [], 4, numeric_toolbox()) + >>> _circles = _packer(4, 0.5, None) + >>> len(_circles) == 4 and feasible(_circles) + True + """ + + @Template.define + def propose_packer_visual( + self, + current: Packer, + feedback: list[Rollout], + n: int, + toolbox: str, + render: Image.Image, + ) -> Packer: + """Write an improved ``pack(n, time_budget, current_best)`` that packs {n} + non-overlapping circles into the unit square, maximizing the SUM OF THE RADII. + + Here is what the current packing actually looks like: + + {render} + + and here is what the evaluator measured: + + + {feedback} + + + Use the picture: wasted space, circles that could grow, and regions that want a + different arrangement are visible in it in a way they are not in the numbers. + + {toolbox} + + Then apply the same rules as before -- read ``n``, spend ``time_budget``, + polish the ``current_best`` you are handed rather than restarting, never return + an infeasible packing -- and put doctests in your docstring certifying that + ``cs = pack(5, 0.5, None)`` yields ``len(cs) == 5 and feasible(cs)``, each + input line prefixed with the doctest prompt (three ``>`` characters and a + space). ``Circle``, ``feasible``, ``total_radius`` and ``worst_violation`` are + in scope. + """ + + @Template.define + def refine_packer( + self, + current: Packer, + instruction: str, + feedback: list[Rollout], + n: int, + toolbox: str, + ) -> Packer: + """Apply a refinement instruction to a packing algorithm. + + This is the *refiner module* being spent: another module of the search evolved + the instruction below, and your job is to carry it out on the current + ``pack(n, time_budget, current_best)`` faithfully -- not to substitute your own + plan for it. + + + {instruction} + + + Here is how the current packer scored, for context on what the instruction is + reacting to: + + + {feedback} + + + {toolbox} + + Return the revised packer for n={n}. It keeps the same contract: read ``n``, + spend ``time_budget``, polish ``current_best`` rather than discarding it, never + return an infeasible packing, and carry doctests in the docstring certifying + that ``cs = pack(5, 0.5, None)`` yields ``len(cs) == 5 and feasible(cs)`` (each + input line prefixed with the doctest prompt -- three ``>`` characters and a + space). If the packer you were given is broken, repair it: recovering a failed + mutation is exactly what this module is for. ``Circle``, ``feasible``, + ``total_radius`` and ``worst_violation`` are in scope. + """ + + @Template.define + def propose_refiner( + self, current: str, packer: Packer, feedback: list[Rollout] + ) -> str: + """You are optimizing the REFINER INSTRUCTION -- the second module of this + search. It is a short natural-language directive that another model applies to + the current packing algorithm to produce the next one, so it is where a + *strategy* can be discovered and held even while the code lags behind it. + + + {current} + + + The algorithm it will be applied to is above, and here is how that algorithm + scored: + + + {feedback} + + + Write a better instruction. It should name the specific structural change worth + making next -- switch how the radii are solved for, change how the centres + move, add a different seeding strategy, escape a saturated configuration, or + repair a broken implementation -- in enough detail that a competent programmer + could carry it out without guessing, while still being an instruction rather + than the code itself. Aim past the current implementation: this module is + valuable precisely when it is ahead of the code. + + Return the instruction as plain text, nothing else. + """ + + @Template.define + def bootstrap_packer(self, objective: str, n: int, toolbox: str) -> Packer: + """Seedless mode: there is no artifact yet, only a goal. + + + {objective} + + + {toolbox} + + Write the first version of ``pack(n, time_budget, current_best)`` for n={n}: it + returns a list of ``Circle(x, y, r)`` filling the unit square without overlaps. + Read ``n``, spend ``time_budget``, and start from ``current_best`` when it is + not ``None``. Your docstring MUST contain doctests certifying that + ``cs = pack(5, 0.5, None)`` yields ``len(cs) == 5 and feasible(cs)``, each + input line prefixed with the doctest prompt (three ``>`` characters and a + space). ``Circle``, ``feasible`` and ``total_radius`` are in scope. + """ + + +# --------------------------------------------------------------------------- +# The task and its evaluator -- deterministic Python, the ground truth every +# candidate is scored by. +# --------------------------------------------------------------------------- + + +class _Timeout(Exception): + pass + + +def _run_packer( + packer: Packer, n: int, time_budget: float, current_best: list[Circle] | None +) -> tuple[list[Circle], float]: + """Run a synthesized packer under a hard wall-clock backstop, returning its + circles and how long it took. The artifact is *given* its budget and the incumbent + packing, and is expected to honour both; the alarm only catches one that does not. + """ + + def _alarm(signum: int, frame: object) -> None: + raise _Timeout(f"pack() ignored its {time_budget}s budget") + + guarded = threading.current_thread() is threading.main_thread() + if guarded: + previous = signal.signal(signal.SIGALRM, _alarm) + signal.setitimer(signal.ITIMER_REAL, time_budget * 3.0 + 5.0) + start = time.perf_counter() + try: + circles = list(packer(n, time_budget, current_best)) + finally: + elapsed = time.perf_counter() - start + if guarded: + signal.setitimer(signal.ITIMER_REAL, 0.0) + signal.signal(signal.SIGALRM, previous) + return circles, elapsed + + +def packing_score(circles: collections.abc.Sequence[Circle], n: int) -> float: + """The one number: the sum of radii after shrinking the packing to exact + feasibility, and zero for a packing of the wrong size. + + >>> packing_score([Circle(0.25, 0.25, 0.25), Circle(0.75, 0.75, 0.25)], 2) + 0.5 + >>> packing_score([Circle(0.5, 0.5, 1.0)], 1) + 0.5 + """ + if len(circles) != n: + return 0.0 + return feasible_scale(circles) * total_radius(circles) + + +def generality_check(packer: Packer, n: int) -> Diagnostic: + """Side information only, never scored: how the packer does on an instance size it + was not asked about. A genuine algorithm keeps its edge here; a table of + coordinates for one ``n`` falls back to whatever it does by default, and the + proposer gets to see that it did. + + Read the verdict with the 0.5s budget in mind. It is short because this is the one + diagnostic that costs an extra run of the packer, and it is short enough that an + artifact which returns a safe fallback rather than a half-finished optimization when + rushed is indistinguishable here from one that memorised an answer -- so a "table" + verdict is evidence about the packer's behaviour under a tight budget, not proof + that it fails to generalize. + """ + other = n + 3 + baseline = total_radius(seed_packer(other, 0.1, None)) + try: + circles, _ = _run_packer(packer, other, 0.5, None) + except Exception as exc: + return Diagnostic( + "generality", f"pack({other}, ...) raised {type(exc).__name__}: {exc}" + ) + if len(circles) != other: + return Diagnostic( + "generality", + f"pack({other}, ...) returned {len(circles)} circles, not {other}", + ) + achieved = packing_score(circles, other) + return Diagnostic( + "generality", + f"on the unrequested size n={other} this packer scores {achieved:.6f} vs " + f"{baseline:.6f} for the naive grid -- " + + ( + "it generalizes" + if achieved > baseline + else "no better than the grid, so it is a table rather than an algorithm" + ), + ) + + +def trajectory_metrics(scores: list[float]) -> list[Metric]: + """The Pareto objectives of the paper's single-task search. + + Its Mechanism 3 says the front is kept across "max score, mean score, EMA + stability, improvement rate", and that this is what keeps greedy, LP, SLP, bilevel + L-BFGS and CMA-ES candidates alive at once. Those are properties of a *run + distribution*, not of one packing, so the evaluator runs each packer several times + and reports the shape of the result: + + * ``max_score`` -- the best packing it found, the headline number + * ``mean_score`` -- what it achieves typically, not at its luckiest + * ``worst_score`` -- the floor it is guaranteed not to fall below + + Three, not the paper's four. It names its four without defining them, and the two + obvious readings of the missing pair do not survive contact with this evaluator: + + * "Improvement rate" as best-minus-first over the repeats measures which draw came + out best, because the repeats are *independent* runs of the same artifact on the + same input rather than a sequence of refinements. It is a property of the random + seed, and no arrangement of it can say what it wants to say -- how much the + artifact would gain from more time -- without actually giving it more time. + * "Stability" as any scale-free measure of run-to-run agreement is maximized by an + artifact that reliably does nothing. A deterministic packer takes the best + attainable value whatever it scores, so it is non-dominated on that axis + permanently and can never be pruned. An objective a do-nothing candidate wins + outright does not keep algorithmic families alive on the frontier; it keeps junk + alive on it. + + ``worst_score`` is the consistency objective that is not gameable that way: it + rewards an artifact for being reliable *at a good level*, and a candidate that + reliably scores nothing is last on it rather than first. All three are + higher-is-better, which the Pareto machinery requires. + + >>> [str(m) for m in trajectory_metrics([1.0, 1.0, 1.0])] + ['max_score=1', 'mean_score=1', 'worst_score=1'] + >>> [str(m) for m in trajectory_metrics([0.4, 1.0])] + ['max_score=1', 'mean_score=0.7', 'worst_score=0.4'] + """ + if not scores: + return [ + Metric("max_score", 0.0), + Metric("mean_score", 0.0), + Metric("worst_score", 0.0), + ] + return [ + Metric("max_score", max(scores)), + Metric("mean_score", statistics.fmean(scores)), + Metric("worst_score", min(scores)), + ] + + +PACKING_REPEATS = 3 + + +def evaluate_packing( + packer: Packer, + n: int, + time_budget: float, + current_best: list[Circle] | None, + *, + diagnose: bool = True, +) -> tuple[Evaluation, list[Circle]]: + """Score a packer and explain the score. + + Ground truth, deterministic Python: the score is the sum of radii after shrinking + the packing to exact feasibility (see ``feasible_scale``), so overlap is paid for + proportionally rather than at a cliff and there is no tolerance to exploit. + + The packer is run ``PACKING_REPEATS`` times, each time handed the incumbent, and + the run distribution becomes the Pareto objectives (``trajectory_metrics``). The + repeats are not redundancy: these artifacts use randomised restarts, so "how good + is it typically" and "does it stay there" are different questions from "how good + was its best run", and the paper keeps candidates that win any of them. + + Everything the evaluator learns on the way -- which circles are jammed, where the + slack is, whether the radii are suspiciously uniform, the spread across repeats, + how it fares on a size it was not asked about, and the traceback if it crashed -- + goes back as Side Information. The best packing found is returned alongside the + evaluation so the caller can make it the incumbent without paying for another run. + + ``diagnose=False`` skips the generality check, which is the one diagnostic that + costs a whole extra run of the packer. The SI ablation passes it: on a domain whose + budget is wall-clock, an arm whose diagnostics are discarded unread must not be + charged for producing them, or the ablation measures the bill as well as the effect. + """ + scores: list[float] = [] + best: list[Circle] = [] + elapsed = 0.0 + for _ in range(PACKING_REPEATS): + try: + circles, seconds = _run_packer(packer, n, time_budget, current_best) + except Exception: + return Evaluation( + score=0.0, + metrics=trajectory_metrics([0.0]), + diagnostics=[ + Diagnostic("crash", traceback.format_exc(limit=3).strip()), + Diagnostic( + "fix", + "pack(n, time_budget, current_best) must return a list of " + "Circle without raising", + ), + ], + ), [] + elapsed = max(elapsed, seconds) + scores.append(packing_score(circles, n)) + if scores[-1] >= max(scores): + best = circles + + metrics = trajectory_metrics(scores) + score = max(scores) + diagnostics: list[Diagnostic] = [ + Diagnostic( + "repeats", + f"{PACKING_REPEATS} runs scored " + + ", ".join(f"{s:.6f}" for s in scores) + + f"; the score is the best of them ({score:.6f})", + ), + Diagnostic("runtime", f"{elapsed:.2f}s of a {time_budget:.2f}s budget per run"), + ] + if current_best is not None: + incumbent = packing_score(current_best, n) + diagnostics.append( + Diagnostic( + "incumbent", + f"the packing handed to you scored {incumbent:.6f}; this artifact " + + ( + f"improved it by {score - incumbent:.6f}" + if score > incumbent + else "did not improve on it, which means the time went nowhere -- " + "start from what you are given" + ), + ) + ) + + if len(best) != n: + return Evaluation( + score=0.0, + metrics=metrics, + diagnostics=diagnostics + + [Diagnostic("count", f"returned {len(best)} circles, expected {n}")], + ), [] + + violation = worst_violation(best) + total = total_radius(best) + scale = feasible_scale(best) + smallest = min(c.r for c in best) + if diagnose: + diagnostics.append(generality_check(packer, n)) + + if violation > 1e-9: + offenders = sorted(best, key=lambda c: -c.r)[:3] + diagnostics += [ + Diagnostic( + "infeasible", + f"worst constraint violation {violation:.6f} (overlap depth or " + f"excursion outside the unit square). Radii sum to {total:.6f} as " + f"returned, but every radius has to shrink by a factor of " + f"{scale:.6f} before the packing is legal, so the score is " + f"{score:.6f}. Place centres so the radii need no shrinking.", + ), + Diagnostic("largest circles", ", ".join(str(c) for c in offenders)), + ] + return Evaluation(score=score, metrics=metrics, diagnostics=diagnostics), best + + # Feasible: report where the slack is, so the proposer knows what to grow. + slacks = [] + for i, a in enumerate(best): + gap = min( + [a.x - a.r, a.y - a.r, 1.0 - a.x - a.r, 1.0 - a.y - a.r] + + [ + math.hypot(a.x - b.x, a.y - b.y) - a.r - b.r + for j, b in enumerate(best) + if j != i + ] + ) + slacks.append((gap, i, a)) + loosest = sorted(slacks, reverse=True)[:3] + tightest = sorted(slacks)[:3] + diagnostics += [ + Diagnostic( + "room to grow", + "; ".join( + f"circle {i} {c} has {gap:.4f} of free space" for gap, i, c in loosest + ) + or "every circle is jammed", + ), + Diagnostic( + "jammed circles", + "; ".join(f"circle {i} {c} slack {gap:.6f}" for gap, i, c in tightest), + ), + Diagnostic( + "radius spread", + f"largest {max(c.r for c in best):.4f}, smallest {smallest:.4f} -- " + + ( + "nearly uniform, which is usually suboptimal" + if max(c.r for c in best) - smallest < 0.01 + else "non-uniform" + ), + ), + ] + return Evaluation(score=score, metrics=metrics, diagnostics=diagnostics), best + + +def seed_packer( + n: int, time_budget: float, current_best: list[Circle] | None +) -> list[Circle]: + """The naive baseline every packing run starts from: equal circles on the + tightest square grid that fits them, unless it is handed something better. + + >>> cs = seed_packer(4, 0.1, None) + >>> len(cs) == 4 and feasible(cs) + True + """ + side = math.ceil(math.sqrt(n)) + r = 1.0 / (2 * side) + grid = [ + Circle(x=(i % side) * 2 * r + r, y=(i // side) * 2 * r + r, r=r) + for i in range(n) + ] + if current_best is not None and packing_score(current_best, n) > packing_score( + grid, n + ): + return list(current_best) + return grid + + +def render_packing(circles: collections.abc.Sequence[Circle], n: int) -> Image.Image: + """Render a packing as a PNG, so Side Information can be visual. Nothing about the + loop changes: an image is simply another Encodable the proposer's prompt splices.""" + import matplotlib + + matplotlib.use("Agg") + from matplotlib import patches, pyplot + + figure = pyplot.figure(figsize=(4, 4), dpi=110) + axes = figure.add_subplot(111, aspect="equal") + axes.add_patch(patches.Rectangle((0, 0), 1, 1, fill=False, linewidth=1.5)) + for i, c in enumerate(circles): + axes.add_patch(patches.Circle((c.x, c.y), c.r, alpha=0.45)) + axes.annotate(str(i), (c.x, c.y), ha="center", va="center", fontsize=7) + axes.set_xlim(-0.05, 1.05) + axes.set_ylim(-0.05, 1.05) + axes.set_title(f"n={n} sum of radii = {total_radius(circles):.4f}") + buffer = io.BytesIO() + figure.savefig(buffer, format="png", bbox_inches="tight") + pyplot.close(figure) + buffer.seek(0) + return Image.open(buffer) + + +# --------------------------------------------------------------------------- +# The control the paper does not run. +# --------------------------------------------------------------------------- + + +def baseline_packing( + n: int, seconds: float, rng: random.Random +) -> tuple[list[Circle], int]: + """Warm-started random-restart SLSQP with no model anywhere in the loop. + + This is the comparison 5.3 is missing, and it is the one that decides what the + headline number means. The paper's only comparators on this instance are other + LLM-driven program-search systems -- AlphaEvolve at 2.635, OpenEvolve at 2.6307 -- + so nothing in it separates "reflective search found a good algorithm" from + "reflective search wrote a competent call to a constrained optimizer". Those are + very different claims, and on an instance whose published values sit within a few + thousandths of each other the difference is the whole result. + + So this function is a fair opponent rather than a straw man: it builds the same + ``4n`` wall constraints and ``n(n-1)/2`` separation constraints over the same ``3n`` + variables that the search's winning artifacts converge on, supplies the analytic + constraint Jacobian, and spends its budget on restarts -- half from random centres, + half warm-started with jitter from its own incumbent, which is the same advantage + `run_pack` gives the artifacts through ``current_best``. What it does not have is a + model choosing what to try next. Whatever margin the search shows over this is the + part attributable to reflection. + + Returns the best packing found and how many restarts fitted in the budget. + """ + import numpy as np + from scipy.optimize import minimize + + upper, lower = np.triu_indices(n, k=1) + rows = np.arange(len(upper)) + eye, zero = np.eye(n), np.zeros((n, n)) + # d(wall constraints)/d(x, y, r): constant, so it is built once. + walls_jac = np.vstack( + [ + np.hstack([eye, zero, -eye]), # x - r >= 0 + np.hstack([zero, eye, -eye]), # y - r >= 0 + np.hstack([-eye, zero, -eye]), # 1 - x - r >= 0 + np.hstack([zero, -eye, -eye]), # 1 - y - r >= 0 + ] + ) + + def constraints(v: typing.Any) -> typing.Any: + x, y, r = v[:n], v[n : 2 * n], v[2 * n :] + walls = np.concatenate([x - r, y - r, 1.0 - x - r, 1.0 - y - r]) + gap = np.hypot(x[upper] - x[lower], y[upper] - y[lower]) - r[upper] - r[lower] + return np.concatenate([walls, gap]) + + def constraints_jac(v: typing.Any) -> typing.Any: + x, y = v[:n], v[n : 2 * n] # the separation gradient does not involve r + dx, dy = x[upper] - x[lower], y[upper] - y[lower] + distance = np.maximum(np.hypot(dx, dy), 1e-12) + pairs = np.zeros((len(upper), 3 * n)) + pairs[rows, upper], pairs[rows, lower] = dx / distance, -dx / distance + pairs[rows, n + upper], pairs[rows, n + lower] = dy / distance, -dy / distance + pairs[rows, 2 * n + upper] = pairs[rows, 2 * n + lower] = -1.0 + return np.vstack([walls_jac, pairs]) + + gradient = np.concatenate([np.zeros(2 * n), -np.ones(n)]) + bounds = [(0.0, 1.0)] * (2 * n) + [(0.0, 0.5)] * n + deadline = time.monotonic() + seconds + best: list[Circle] = [] + restarts = 0 + while time.monotonic() < deadline: + restarts += 1 + if best and restarts % 2 == 0: # polish the incumbent + start = np.array( + [c.x for c in best] + [c.y for c in best] + [c.r for c in best] + ) + start[: 2 * n] += np.array([rng.gauss(0.0, 0.02) for _ in range(2 * n)]) + else: # a fresh configuration + start = np.array( + [rng.random() for _ in range(2 * n)] + + [0.5 / math.ceil(math.sqrt(n))] * n + ) + np.clip(start, 0.0, 1.0, out=start) + try: + solved = minimize( + lambda v: -v[2 * n :].sum(), + start, + jac=lambda _: gradient, + method="SLSQP", + bounds=bounds, + constraints=[ + {"type": "ineq", "fun": constraints, "jac": constraints_jac} + ], + options={"maxiter": 200, "ftol": 1e-10}, + ) + except Exception: # a restart that fails to converge is simply skipped + continue + found = [ + Circle( + x=float(solved.x[i]), + y=float(solved.x[n + i]), + r=float(solved.x[2 * n + i]), + ) + for i in range(n) + ] + if packing_score(found, n) > packing_score(best, n): + best = found + scale = feasible_scale(best) + return [Circle(x=c.x, y=c.y, r=c.r * scale) for c in best], restarts + + +def baseline_note(n: int, seconds: float, rng: random.Random) -> str: + """Run the no-LLM control and say what it means, or say why it could not run.""" + if importlib.util.find_spec("scipy.optimize") is None: + return ( + "No-LLM baseline: not available, because scipy is not installed here. The " + "search's numbers are therefore unattributed -- there is nothing to say how " + "much of the distance from the seed is the reflection and how much is the " + "optimizer the artifacts call." + ) + started = time.monotonic() + circles, restarts = baseline_packing(n, seconds, rng) + return ( + f"No-LLM baseline (warm-started random-restart SLSQP, no model in the loop): " + f"{packing_score(circles, n):.7f} for n={n} in " + f"{time.monotonic() - started:.0f}s over {restarts} restarts. This is the " + f"comparison the paper's 5.3 does not report, and the number the search has to " + f"beat for its margin to be about reflection rather than about scipy." + ) + + +# --------------------------------------------------------------------------- +# Wiring the domain to the engine. +# --------------------------------------------------------------------------- + + +def run_pack(args: argparse.Namespace, rng: random.Random) -> tuple[Result, list[str]]: + """Single-task search over a two-module system, with the incumbent threaded through. + + Three things here are the paper's setup rather than a simplification of it: the + artifact is handed the best packing found so far, the Pareto objectives are the + run-distribution metrics of its Mechanism 3, and the search alternates between two + modules on one shared frontier (Mechanism 2). The incumbent is ordinary mutable + state in this closure -- and because it makes an evaluation depend on more than the + artifact, it is declared to the engine as a ``state_key``. + + Returns the result and the per-iteration module choices, which the engine has no + reason to know about and `module_note` needs in order to say which module each + accepted gain came from. + """ + n, budget = args.num_circles, args.time_budget + toolbox = numeric_toolbox() + incumbent: list[list[Circle]] = [[]] # a one-slot cell: the best packing so far + origins: list[str] = [] + + def state_key() -> str: + return f"{packing_score(incumbent[0], n):.12f}" + + def evaluator(system: PackSystem, _: None) -> Evaluation: + evaluation, best = evaluate_packing( + system.packer, + n, + budget, + incumbent[0] or None, + diagnose=not args.no_side_info, + ) + # Whatever the candidate managed becomes the incumbent for everything after it, + # so later artifacts start where this one stopped -- the paper's + # ``current_best_solution``, which is why this domain declares a ``state_key``. + if packing_score(best, n) > packing_score(incumbent[0], n): + incumbent[0] = best + return evaluation + + def mutate_code(system: PackSystem, feedback: list[Rollout]) -> PackSystem: + if not args.visual_si: + packer = Proposer().propose_packer(system.packer, feedback, n, toolbox) + else: + try: + circles, _ = _run_packer(system.packer, n, budget, incumbent[0] or None) + packer = Proposer().propose_packer_visual( + system.packer, feedback, n, toolbox, render_packing(circles, n) + ) + except Exception: # a packer that crashes has nothing to show + packer = Proposer().propose_packer(system.packer, feedback, n, toolbox) + return PackSystem(packer=packer, refiner=system.refiner, origin="code") + + def mutate_refiner(system: PackSystem, feedback: list[Rollout]) -> PackSystem: + # The refiner module advances in two steps: rewrite the instruction, then spend + # it on the current code. The instruction can therefore describe a strategy the + # code does not implement yet -- which is exactly the leapfrogging the paper + # describes, and also how a broken packer gets repaired instead of abandoned. + instruction = Proposer().propose_refiner( + system.refiner, system.packer, feedback + ) + packer = Proposer().refine_packer( + system.packer, instruction, feedback, n, toolbox + ) + return PackSystem(packer=packer, refiner=instruction, origin="refiner") + + def proposer(system: PackSystem, feedback: list[Rollout]) -> PackSystem: + code = rng.random() < args.code_share + # Recorded before the call, so the list stays aligned with the trace even when + # the proposal dies at decode: the engine appends a Step either way. + origins.append("code" if code else "refiner") + return (mutate_code if code else mutate_refiner)(system, feedback) + + def bootstrap(objective: str) -> PackSystem: + return PackSystem( + packer=Proposer().bootstrap_packer(objective, n, toolbox), + refiner=SEED_REFINER, + origin="bootstrap", + ) + + # Annotated rather than inferred: this domain has no dataset, so the element type + # would infer as ``None`` and fail the engine's ``E: Example`` bound. + result: Result = optimize_anything( + evaluator=evaluator, + proposer=proposer, + seed=( + None + if args.seedless + else PackSystem(packer=seed_packer, refiner=SEED_REFINER) + ), + bootstrap=bootstrap, + objective=( + f"Pack {n} non-overlapping circles into the unit square so that the sum " + f"of their radii is as large as possible." + ), + budget=args.budget, + selection=args.selection, + use_side_info=not args.no_side_info, + rng=rng, + task_name=f"pack-{n}", + state_key=state_key, + ) + return result, origins + + +# --------------------------------------------------------------------------- +# Reporting and main +# --------------------------------------------------------------------------- + + +def module_note(result: Result, origins: list[str]) -> str: + """Which module did the work, as far as it can honestly be attributed. + + Multi-module search is only visible if you count it, and counting the survivors on + the frontier is not enough: it says which module's proposals lasted, not which one + moved the score. So this also reports, per module, how many proposals it made, how + many were accepted, and the mean minibatch gain when they were -- and the order the + accepted gains arrived in, which is where the paper's leapfrogging would show up as + the two modules handing off to each other. + + What it deliberately does not report is a per-module score, because there isn't one + (see `PackSystem`). A handover in the sequence below means the other module produced + the next accepted gain; it does not mean that module was ahead of its partner. + """ + surviving = collections.Counter( + c.artifact.origin for c in result.pool if isinstance(c.artifact, PackSystem) + ) + proposed: collections.Counter[str] = collections.Counter() + accepted: collections.Counter[str] = collections.Counter() + gains: dict[str, list[float]] = {"code": [], "refiner": []} + sequence: list[str] = [] + for step, origin in zip(result.history, origins): + proposed[origin] += 1 + if step.accepted: + accepted[origin] += 1 + gains[origin].append(step.after - step.before) + sequence.append(origin) + + winner = result.best.artifact + handovers = sum(a != b for a, b in zip(sequence, sequence[1:])) + return "\n".join( + [ + "Modules: " + + "; ".join( + f"{name} proposed {proposed[name]}, accepted {accepted[name]}" + + ( + f", mean minibatch gain {statistics.fmean(gains[name]):+.6f}" + if gains[name] + else "" + ) + for name in ("code", "refiner") + if proposed[name] + ), + "Surviving frontier by module: " + + (", ".join(f"{k} {v}" for k, v in sorted(surviving.items())) or "none"), + "The best candidate came from the " + + f"{winner.origin if isinstance(winner, PackSystem) else 'unknown'} module", + f"Accepted gains in order: {' -> '.join(sequence) or 'none'}" + + (f" ({handovers} handover(s))" if len(sequence) > 1 else ""), + ] + ) + + +def cold_start_note(result: Result, n: int, time_budget: float) -> str: + """What the winning artifact scores on its own, with no incumbent to polish. + + The headline of a run with the incumbent threaded through belongs to the *run*, not + to the artifact credited with it. Every candidate is handed the best packing found so + far and told never to return something worse, so its score accumulates the work of + everything that ran before it: a candidate that returned its input unchanged would be + recorded at the full incumbent value. That is the paper's own setup, not a deviation + from it -- its evolved packer takes ``current_best_solution`` too, so its 2.63598 is + a trajectory number in exactly the same way -- but it means "the winning artifact + reached X" is not a statement this search is entitled to make. Running the winner + once from nothing is, and it costs one evaluation. + """ + system = result.best.artifact + if not isinstance(system, PackSystem): + return "" + evaluation, _ = evaluate_packing( + system.packer, n, time_budget, None, diagnose=False + ) + return ( + f"Cold start: the winning artifact alone, handed no incumbent, scores " + f"{evaluation.score:.7f} against the run's {result.best_score:.7f}. The " + f"difference is what it inherited from the candidates before it rather than " + f"earned -- the run's number is the search's, the cold-start number is the " + f"artifact's." + ) + + +def render_system(artifact: typing.Any) -> str: + """Both modules of the winning system: the instruction, then the code.""" + if not isinstance(artifact, PackSystem): + return str(artifact) + packer = source_of(artifact.packer) or repr(artifact.packer) + return ( + f"--- refiner instruction ---\n{artifact.refiner}\n\n" + f"--- packer ---\n{packer.rstrip()}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--num-circles", + type=int, + default=10, + help="Circles to pack; the paper's instance is 26", + ) + parser.add_argument("--budget", type=int, default=8, help="Optimizer iterations") + parser.add_argument( + "--time-budget", + type=float, + default=2.0, + help="Seconds a synthesized packer is given per call", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Seed for selection and module choice" + ) + parser.add_argument( + "--code-share", + type=float, + default=0.5, + help="Fraction of iterations that mutate the code module rather than the " + "refiner module (both live on the one shared frontier)", + ) + parser.add_argument( + "--selection", + choices=["pareto", "best"], + default="pareto", + help="Candidate selection; 'best' mutates the best average instead, which is " + "the naive alternative the paper's 4.3 argues against rather than an ablation " + "it runs", + ) + parser.add_argument( + "--no-side-info", + action="store_true", + help="Score-only feedback: the paper's SI ablation", + ) + parser.add_argument( + "--visual-si", + action="store_true", + help="Send a rendered image of the packing as side information", + ) + parser.add_argument( + "--seedless", + action="store_true", + help="Bootstrap candidate zero from a natural-language objective", + ) + parser.add_argument( + "--baseline", + action="store_true", + help="Run the no-LLM control instead of the search: random-restart SLSQP for " + "the wall-clock the search would have spent on packers", + ) + parser.add_argument( + "--baseline-seconds", + type=float, + default=0.0, + help="Seconds for --baseline; 0 matches the search's packer time, which is " + "--budget x --time-budget x the evaluator's repeats", + ) + args = parser.parse_args() + + if args.baseline: + seconds = ( + args.baseline_seconds or args.budget * PACKING_REPEATS * args.time_budget + ) + print( + f"[baseline] no LLM, n={args.num_circles}, {seconds:.0f}s " + + ( + "(as given)" + if args.baseline_seconds + else f"(= {args.budget} iterations x {PACKING_REPEATS} repeats x " + f"{args.time_budget:.0f}s, the packer time the search would spend)" + ) + ) + print(baseline_note(args.num_circles, seconds, random.Random(args.seed))) + return + + result, origins = run_pack(args, random.Random(args.seed)) + report( + result, + selection=args.selection, + side_info=not args.no_side_info, + notes=[ + note + for note in ( + module_note(result, origins), + cold_start_note(result, args.num_circles, args.time_budget), + ) + if note + ], + render_artifact=render_system, + ) + # No assertion that the score improved: one would look reassuring and could not + # fail. The headline is a maximum over the surviving pool, the seed is a candidate + # in it, and a pruned seed was by definition dominated by a survivor. The numbers + # worth checking are the two the notes above print -- the winner's cold-start score, + # and what --baseline reaches with no model at all. + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/prompting.py b/docs/source/llm_examples/optimization/prompting.py new file mode 100644 index 000000000..e2b609f10 --- /dev/null +++ b/docs/source/llm_examples/optimization/prompting.py @@ -0,0 +1,368 @@ +"""Prompt optimization: generalization mode (optimize_anything A.3). + +Optimize a system prompt so it works on instances the search never saw. Both a training +set and a validation set are supplied, which is what selects the paper's generalization +mode: search takes its feedback from the training instances, and the artifact that +survives has to carry to held-out ones. This is the mode GEPA and MIPROv2 operate in, +and the one the paper extends beyond prompts. + +The artifact is a plain ``str``, the evaluator is itself a model call, and -- as in the +paper -- the prompt is optimized *for a cheaper model* (``--worker-model``) than the one +proposing it, here GPT-4.1-mini, which is A.3's target model too. In effectful "run this +call on a different model" is a scoped handler, so `library.worker` is the entire +mechanism. + +The task is constrained writing -- an exact word count, an initial-letter rule, a banned +letter -- scored by deterministic Python, and that substitution needs stating plainly +rather than defending. AIME itself would have been the faithful choice and has ample +headroom: the paper measures GPT-4.1-mini at 46.67% on AIME 2025 from a generic prompt. +What ruled it out is that the problems cannot be embedded here -- the set is large, and +writing a dozen substitutes is not AIME. A substitute set also saturates: GPT-4.1-mini +scores 12/12 on hand-written counting and number-theory problems from the bare seed +prompt, which leaves the search nothing to climb. Constraint tracking is a task +these models do fail, the checker is exact, and the lever a better prompt supplies is +method -- count before answering, verify each constraint separately, revise once. +Nothing measured here transfers to a claim about AIME. + +Demonstrates: +- Generalization mode: per-example Pareto objectives over the training instances, so a + prompt that is best at *something* survives, and selection on a held-out set +- Side Information following the paper's design for this domain as far as the task + allows -- the instance, the model's reasoning, what it produced, and a per-constraint + account of what went wrong. A.3 also returns the ground-truth answer, which has no + analogue here: constrained writing has no reference sentence, only constraints +- The proposer/target-model split as a scoped handler nested inside the harness's stack +- Partial credit as a search gradient: a 0/1 verdict would make most of the run invisible + +Measured on 2026-07-29 with gpt-5.5 proposing and gpt-4.1-mini writing, 8 iterations: +training score 0.733 -> 0.867, held-out 0.800 -> 0.800. The search improved the prompt +on the instances it saw and none of that carried, which is left standing rather than +tuned away -- but it is a weak observation, not a negative result. Five validation +instances scored in thirds resolve nothing below about 0.07, one run gives no variance +estimate, and the held-out score is also the set the winner was selected on. It says +this run did not show transfer, and no more than that. +""" + +# Simplifications vs. the source: +# - The task is constrained writing, not AIME 2022-2025, for the reason above. +# - Five training and five validation instances. The paper trains on AIME 2022-2024 +# and tests on AIME 2025 -- on the order of a hundred problems and thirty, with +# Figure 7's 57.78% validation score implying a 45-problem validation split. +# - The winner is selected on the same held-out set this script then reports, where the +# paper keeps a third split and reports the test score. Read the val number as +# selection-biased. +# - Budget is counted in optimizer iterations, not metric calls or dollars: about 60 +# evaluator calls here against the paper's ~350 and $6.44. `report` prints the count. +# - There is no baseline optimizer. A.3's actual claim is 60.0% against MIPROv2's +# 51.33% on the same benchmark; nothing here compares against any other optimizer, +# or even against best-of-N hand-written prompts, so that claim is untouched. +# - The candidate is spliced ahead of the question in a user message rather than being +# a system prompt, and the answer comes back as a typed ``Answer(reasoning, final)``. +# The type therefore supplies two of the things the paper's evolved prompt had to +# learn -- an explicit reasoning step, and isolating the final answer (its rule 6) -- +# so roughly a third of Appendix J's content is unreachable as a lever here. +# - One run, one sample per instance, frozen thereafter by the evaluation cache: no +# repeats, no seed sweep, no variance estimate. +# - Every score is for the prompt *plus the harness's retry loop*, not for the prompt +# alone. ``worker(...)`` scopes the model but does not shadow the ``TenacityRetryer`` +# above it, so an answer that fails to decode is fed its own error and asked again; +# only exhausting the retries reaches the ``except`` here and scores zero. A prompt +# whose answers are borderline-undecodable is flattered by that. +# - The winner is a maximum over the frontier's validation scores while the seed is a +# single validation measurement, so the reported delta is biased upward. Only +# frontier candidates are validated at all, so a candidate that the training set +# dominates can never be selected however well it generalizes. +# - One proposer model; the paper also reports a weaker-proposer arm (its Table 8). + +import argparse +import random +import traceback + +import pydantic.dataclasses + +from docs.source.llm_examples.optimization.library import ( + WORKER_MODEL, + Diagnostic, + Evaluation, + Result, + Rollout, + optimize_anything, + report, + worker, +) +from effectful.handlers.llm import Agent, Template + + +@pydantic.dataclasses.dataclass(frozen=True) +class Writing: + """One constrained-writing instance. The constraints are in the question, so + nothing is hidden from the answering model: what a better prompt supplies is the + *method* for satisfying them, which is exactly what the paper's optimized prompts + encode.""" + + name: str + topic: str + words: int + initial: str + banned: str + + @property + def question(self) -> str: + return ( + f"Write a single sentence about {self.topic}. It must contain exactly " + f"{self.words} words, every word must begin with the letter " + f"'{self.initial}', and the letter '{self.banned}' must not appear " + f"anywhere in the sentence." + ) + + +WRITINGS: list[Writing] = [ + Writing("sailing", "sailing", 6, "s", "e"), + Writing("markets", "morning markets", 7, "m", "a"), + Writing("cats", "curious cats", 5, "c", "i"), + Writing("planets", "distant planets", 8, "p", "o"), + Writing("bridges", "old bridges", 6, "b", "u"), + Writing("trains", "night trains", 7, "t", "e"), + Writing("gardens", "walled gardens", 5, "g", "s"), + Writing("harbours", "quiet harbours", 6, "h", "i"), + Writing("lanterns", "paper lanterns", 7, "l", "o"), + Writing("rivers", "wide rivers", 5, "r", "a"), +] + +TRAIN = [ + w + for w in WRITINGS + if w.name in {"sailing", "markets", "cats", "planets", "bridges"} +] +VAL = [w for w in WRITINGS if w not in TRAIN] + +SEED_PROMPT = "Answer the question." + + +@pydantic.dataclasses.dataclass(frozen=True) +class Answer: + """What the answering model returns: its reasoning and its final answer.""" + + reasoning: str + final: str + + +@Template.define +def answer_question(instructions: str, question: str) -> Answer: + """{instructions} + + + {question} + + """ + + +def words_of(sentence: str) -> list[str]: + """The sentence's words, stripped of punctuation and lowercased. + + >>> words_of("Silent ships sail; south, softly.") + ['silent', 'ships', 'sail', 'south', 'softly'] + """ + cleaned = "".join( + c if c.isalpha() or c.isspace() or c == "'" else " " for c in sentence + ) + return [w for w in cleaned.lower().split() if w] + + +def score_writing(sentence: str, task: Writing) -> tuple[float, list[Diagnostic]]: + """Check the three constraints and explain every miss. + + Partial credit on purpose: a 0/1 verdict would make most of the search invisible, + while per-constraint credit is a gradient the proposer can climb -- and the + per-constraint breakdown *is* the side information. + + A sentence that satisfies all three constraints of the first instance (six words, + every word starting with 's', no letter 'e' anywhere) scores 1.0: + + >>> score, _ = score_writing("Ships sail south, ships spin swiftly.", WRITINGS[0]) + >>> score + 1.0 + + while the near-miss "Silent ships sail south, softly singing." -- same six words, + same initial, but 'silent' smuggles in an 'e' -- loses exactly one third: + + >>> score, _ = score_writing("Silent ships sail south, softly singing.", WRITINGS[0]) + >>> round(score, 4) + 0.6667 + """ + words = words_of(sentence) + count_ok = len(words) == task.words + starting = [w for w in words if w.startswith(task.initial)] + initial_ratio = len(starting) / len(words) if words else 0.0 + banned_hits = sentence.lower().count(task.banned) + + diagnostics = [ + Diagnostic("sentence", repr(sentence)), + Diagnostic( + "word count", + f"{len(words)} words {words}, needed exactly {task.words}" + if not count_ok + else f"exactly {task.words} words, as required", + ), + Diagnostic( + "initial letter", + f"{len(starting)}/{len(words)} words begin with '{task.initial}'" + + ( + "" + if initial_ratio == 1.0 + else f"; offending words: {[w for w in words if not w.startswith(task.initial)]}" + ), + ), + Diagnostic( + "banned letter", + f"the letter '{task.banned}' appears {banned_hits} time(s) and must not appear" + if banned_hits + else f"the letter '{task.banned}' does not appear, as required", + ), + ] + score = ( + (1.0 if count_ok else 0.0) + initial_ratio + (1.0 if banned_hits == 0 else 0.0) + ) / 3.0 + return score, diagnostics + + +def evaluate_prompt(prompt: str, task: Writing | None, model: str) -> Evaluation: + """Run one writing instance under the candidate prompt and score it. + + The Side Information follows the paper's design for this domain: the instance, the + model's reasoning, what it produced, and a per-constraint account of what went + wrong -- not merely that something did. + """ + assert task is not None, "the prompt domain always has a dataset" + try: + with worker(model): + produced = answer_question(prompt, task.question) + except Exception: + return Evaluation( + score=0.0, + diagnostics=[ + Diagnostic("task", task.question), + Diagnostic("crash", traceback.format_exc(limit=2).strip()), + ], + ) + score, diagnostics = score_writing(produced.final, task) + return Evaluation( + score=score, + diagnostics=[ + Diagnostic("task", task.question), + Diagnostic("reasoning", produced.reasoning), + *diagnostics, + Diagnostic( + "verdict", + f"scored {score:.2f} of 1.00 -- one third for the exact word count, one " + f"third for the fraction of words with the right initial, one third for " + f"avoiding the banned letter", + ), + ], + ) + + +class Proposer(Agent): + """You are a reflective optimizer. You are shown the current prompt, the score it + achieved, and diagnostic side information explaining *why* it scored that way, and + you return a better prompt. You do not mutate blindly: you first read the + diagnostics to decide which failure mode is costing the most, then you write the + instruction that addresses it.""" + + @Template.define + def propose_prompt(self, current: str, feedback: list[Rollout]) -> str: + """You are optimizing the SYSTEM PROMPT given to a small model that writes + sentences under hard constraints -- an exact word count, a required initial + letter for every word, and a letter that must not appear. The prompt below is + the artifact; return an improved one. + + + {current} + + + Here is how it did on a few instances, with the model's own reasoning, the + sentence it produced, and a per-constraint account of what went wrong: + + + {feedback} + + + Diagnose before you rewrite. The constraints are always stated in the task + itself, so the prompt's job is not to repeat them but to supply a *method* + that makes them stick: how to construct the sentence so the count is right by + construction, how to check each constraint separately rather than trusting a + glance, what to do on finding a violation, and which failure the feedback + shows is currently costing the most. Encode that as durable, general + instructions -- the prompt is scored on instances you have not seen, with + different topics, counts, letters and banned letters, so never mention a + specific instance, and never write a sentence yourself. + + Return the improved prompt as plain text, nothing else. + """ + + +# --------------------------------------------------------------------------- +# Wiring and main +# --------------------------------------------------------------------------- + + +def run_prompt(args: argparse.Namespace, rng: random.Random) -> Result: + return optimize_anything( + evaluator=lambda prompt, task: evaluate_prompt(prompt, task, args.worker_model), + proposer=lambda prompt, feedback: Proposer().propose_prompt(prompt, feedback), + seed=SEED_PROMPT, + dataset=TRAIN, + valset=VAL, + budget=args.budget, + minibatch_size=args.minibatch, + selection=args.selection, + use_side_info=not args.no_side_info, + rng=rng, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--budget", type=int, default=8, help="Optimizer iterations") + parser.add_argument( + "--minibatch", type=int, default=2, help="Instances per reflection step" + ) + parser.add_argument( + "--seed", type=int, default=0, help="Seed for selection and minibatches" + ) + parser.add_argument( + "--worker-model", + default=WORKER_MODEL, + help="Model the prompt is optimized FOR; the harness's --model is the proposer, " + "as in the paper's proposer/worker split", + ) + parser.add_argument( + "--selection", + choices=["pareto", "best"], + default="pareto", + help="Candidate selection; 'best' mutates the best average instead, which is " + "the naive alternative the paper's 4.3 argues against rather than an ablation " + "it runs", + ) + parser.add_argument( + "--no-side-info", + action="store_true", + help="Score-only feedback: the paper's SI ablation", + ) + args = parser.parse_args() + + result = run_prompt(args, random.Random(args.seed)) + report(result, selection=args.selection, side_info=not args.no_side_info) + # No assertion that the score improved. In generalization mode the seed can be + # pruned as training-dominated while every surviving candidate is worse on the + # held-out set, and that is a legitimate outcome of a search this small -- the + # result to report, not a failure to raise on. + if result.best_score <= result.seed_score: + print( + "\nThe search did not improve the held-out score. On five validation " + "instances that is as likely to be the budget as the method." + ) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/__init__.py b/docs/source/llm_examples/reasoning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/reasoning/aime2024.py b/docs/source/llm_examples/reasoning/aime2024.py new file mode 100644 index 000000000..bc00b5999 --- /dev/null +++ b/docs/source/llm_examples/reasoning/aime2024.py @@ -0,0 +1,147 @@ +"""Solving hard problems by writing and running Python. + +You are a careful problem solver and an expert Python programmer. You answer by +writing code, not by reasoning in prose alone: problems that are error-prone to +work out by hand are often easy to brute-force or verify with a short program. + +Each template below generalizes a single problem from the 2024 AIME II over +one or more of its constants; passing the original contest constant recovers +the official answer (noted per template). +""" + +import argparse + +from effectful.handlers.llm import Template + + +@Template.define +def least_beautiful_base(threshold: int) -> int: + r"""Find the least integer base b >= 2 for which there are more than + {threshold} ``b``-eautiful integers. + + A positive integer n is ``b``-eautiful if it has exactly two digits when + written in base b and those two digits sum to ``sqrt(n)``. For example, 81 + is 13-eautiful because 81 = 6_3 in base 13 and 6 + 3 = sqrt(81). + + >>> least_beautiful_base(0) + 3 + >>> least_beautiful_base(1) + 7 + >>> least_beautiful_base(5) + 31 + >>> least_beautiful_base(7) + 211 + """ + + +@Template.define +def root_of_unity_product(n: int) -> int: + r"""Let omega != 1 be a primitive n-th root of unity, for n = {n}. Find the + remainder when the product, over k = 0, ..., n - 1, of + (2 - 2 * omega^k + omega^(2k)) is divided by 1000. + + >>> root_of_unity_product(3) + 13 + >>> root_of_unity_product(5) + 41 + >>> root_of_unity_product(7) + 113 + >>> root_of_unity_product(13) + 321 + """ + + +@Template.define +def max_chip_placements(k: int) -> int: + r"""There is a collection of k^2 indistinguishable black chips and k^2 + indistinguishable white chips, for k = {k}. Find the number of ways to + place some of these chips in the k^2 unit cells of a k-by-k grid so that + all chips in the same row and all chips in the same column have the same + color, and any additional chip placed on the grid would violate one or + more of the previous two conditions. + + >>> max_chip_placements(1) + 2 + >>> max_chip_placements(2) + 6 + >>> max_chip_placements(3) + 38 + >>> max_chip_placements(5) + 902 + """ + + +@Template.define +def count_symmetric_triples(n: int, target: int) -> int: + r"""Find the number of triples of nonnegative integers (a, b, c) satisfying + a + b + c = {n} and + a^2*b + a^2*c + b^2*a + b^2*c + c^2*a + c^2*b = {target}. + + >>> count_symmetric_triples(3, 6) + 7 + >>> count_symmetric_triples(6, 48) + 13 + >>> count_symmetric_triples(9, 162) + 19 + >>> count_symmetric_triples(300, 6000000) + 601 + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="problem", required=True) + + p14 = subparsers.add_parser("least-beautiful-base", help="2024 AIME II Problem 14") + p14.add_argument( + "--threshold", + type=int, + default=10, + help="Find the least base with more than this many b-eautiful integers", + ) + + p13 = subparsers.add_parser("root-of-unity-product", help="2024 AIME II Problem 13") + p13.add_argument( + "--n", + type=int, + default=13, + help="Order of the root of unity", + ) + + p9 = subparsers.add_parser("max-chip-placements", help="2024 AIME II Problem 9") + p9.add_argument( + "--k", + type=int, + default=5, + help="Side length of the grid (and number of chips of each color, k^2)", + ) + + p11 = subparsers.add_parser( + "count-symmetric-triples", help="2024 AIME II Problem 11" + ) + p11.add_argument("--n", type=int, default=300, help="Required sum a + b + c") + p11.add_argument( + "--target", + type=int, + default=6_000_000, + help="Required value of a^2 b + a^2 c + b^2 a + b^2 c + c^2 a + c^2 b", + ) + + args = parser.parse_args() + + if args.problem == "least-beautiful-base": + print(f"Least b with > {args.threshold} b-eautiful integers") + print(f"Answer: {least_beautiful_base(args.threshold)}") + elif args.problem == "root-of-unity-product": + print(f"Product over {args.n}-th roots of unity, mod 1000") + print(f"Answer: {root_of_unity_product(args.n)}") + elif args.problem == "max-chip-placements": + print(f"Maximal chip placements on a {args.k}-by-{args.k} grid") + print(f"Answer: {max_chip_placements(args.k)}") + elif args.problem == "count-symmetric-triples": + print(f"Triples with a + b + c = {args.n} and symmetric sum = {args.target}") + print(f"Answer: {count_symmetric_triples(args.n, args.target)}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/constrained_paragraph.py b/docs/source/llm_examples/reasoning/constrained_paragraph.py new file mode 100644 index 000000000..f75102a9d --- /dev/null +++ b/docs/source/llm_examples/reasoning/constrained_paragraph.py @@ -0,0 +1,48 @@ +"""Solving hard problems by writing and running Python. + +You are a careful problem solver and an expert Python programmer. You answer by +writing code, not by reasoning in prose alone: problems that are error-prone to +work out by hand are often easy to brute-force or verify with a short program. +""" + +import argparse + +from effectful.handlers.llm import Template + + +@Template.define +def constrained_paragraph(endings: list[str]) -> str: + r"""Write a short paragraph whose sentences end, in order, with the words in + {endings}: one sentence per word, each ending with that exact word. + + The examples below split the returned paragraph into sentences and compare the + last word of each (lowercased, punctuation stripped) against the requested + endings -- so a synthesized function must build text with the right shape: + + >>> import re + >>> def endings_of(paragraph): + ... sents = [s for s in re.split(r"(?<=[.!?])\s+", paragraph.strip()) if s] + ... return [re.findall(r"[A-Za-z']+", s)[-1].lower() for s in sents] + >>> endings_of(constrained_paragraph(["walk", "tumbling", "another", "lunatic"])) + ['walk', 'tumbling', 'another', 'lunatic'] + >>> endings_of(constrained_paragraph(["dawn", "river"])) + ['dawn', 'river'] + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--endings", + nargs="+", + default=["mountain", "whisper", "thunder"], + metavar="WORD", + help="Words each sentence must end with, in order", + ) + args = parser.parse_args() + print(f"Paragraph with sentences ending in {args.endings}") + print(f"Answer: {constrained_paragraph(args.endings)}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/countdown.py b/docs/source/llm_examples/reasoning/countdown.py new file mode 100644 index 000000000..ed7807e07 --- /dev/null +++ b/docs/source/llm_examples/reasoning/countdown.py @@ -0,0 +1,83 @@ +""" +In-context learning to solve problems with code across a conversation. +""" + +import argparse +import collections.abc + +from effectful.handlers.llm import Agent, Template + + +class CountdownSolver(Agent): + """ + You are a careful problem solver and an expert Python programmer. You answer by + writing code, not by reasoning in prose alone: problems that are error-prone to + work out by hand are often easy to brute-force or verify with a short program. + """ + + @Template.define + def solve(self, numbers: collections.abc.Sequence[int], target: int) -> bool: + """In the Countdown numbers game, decide whether {target} can be made from + {numbers}, using each number exactly once and combining them with + - * / + (every intermediate division must come out exact). + + >>> agent = CountdownSolver() + >>> agent.solve([2, 3, 5], 11) + True + >>> agent.solve([1, 1], 5) + False + >>> agent.solve([4, 7, 8, 9], 100) + True + >>> agent.solve([5, 5, 5], 3) + False + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--numbers", + nargs="+", + type=int, + default=None, + metavar="N", + help="Numbers to combine (used with --target for a single problem)", + ) + parser.add_argument( + "--target", + type=int, + default=None, + help="Target value to make from --numbers", + ) + args = parser.parse_args() + if (args.numbers is None) != (args.target is None): + parser.error("--numbers and --target must be given together") + + agent = CountdownSolver() + + # A custom problem has no known answer to validate against, so just solve it. + if args.numbers is not None: + print(f"Testing solve({args.numbers}, {args.target})...") + answer = agent.solve(args.numbers, args.target) + print(f"solve({args.numbers}, {args.target}): {answer}") + return + + # Fresh examples (none appear in the docstring doctests), each paired with its + # known-correct answer so we can validate the agent's output. + test_examples: list[tuple[list[int], int, bool]] = [ + ([3, 6, 25, 50], 147, True), # (50 - 25) * 6 - 3 + ([1, 2, 3, 4], 24, True), # 1 * 2 * 3 * 4 + ([2, 4, 8], 9, False), # all-even operands can never reach an odd target + ] + for numbers, target, expected in test_examples: + print(f"Testing solve({numbers}, {target})...") + answer = agent.solve(numbers, target) + status = "OK" if answer == expected else "WRONG" + print(f"[{status}] solve({numbers}, {target}): {answer} (expected {expected})") + assert answer == expected, ( + f"solve({numbers}, {target}) = {answer}, expected {expected}" + ) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/fix_typos.py b/docs/source/llm_examples/reasoning/fix_typos.py new file mode 100644 index 000000000..16a966b21 --- /dev/null +++ b/docs/source/llm_examples/reasoning/fix_typos.py @@ -0,0 +1,49 @@ +"""Solving hard problems by writing and running Python. + +You are a careful problem solver and an expert Python programmer. You answer by +writing code, not by reasoning in prose alone: problems that are error-prone to +work out by hand are often easy to brute-force or verify with a short program. +""" + +import argparse + +from effectful.handlers.llm import Template + + +@Template.define +def fix_typos(text: str) -> str: + """Output the following text exactly, with no changes at all except for fixing + the misspellings. Leave every other stylistic decision -- commas, US vs British + spellings, capitalization, line breaks -- exactly as in the original: + + {text} + + Only misspelled words may change; every correctly spelled word and all + punctuation and whitespace must be preserved verbatim. Identify the typos, then + apply the corrections with code so that nothing else can drift. + + >>> fix_typos("We inctroduce a probablistic method in the presense of noise.") + 'We introduce a probabilistic method in the presence of noise.' + >>> fix_typos("Teh quick borwn fox jumpps over the lazy dog.") + 'The quick brown fox jumps over the lazy dog.' + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--text", + type=str, + default=( + "We inctroduce a probablistic algorithm that estimates the " + "timne-varying location in the presense of measurment noise." + ), + help="Text whose typos should be fixed", + ) + args = parser.parse_args() + print(f"Fix only the typos in:\n{args.text}") + print(f"Answer: {fix_typos(args.text)}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/gridworlds.py b/docs/source/llm_examples/reasoning/gridworlds.py new file mode 100644 index 000000000..5461b9664 --- /dev/null +++ b/docs/source/llm_examples/reasoning/gridworlds.py @@ -0,0 +1,140 @@ +import abc +import dataclasses +import enum +import typing + + +@dataclasses.dataclass(frozen=True, eq=True, unsafe_hash=True) +class State: + """A raw grid observation: rows of integer color codes. + + A ``dataclass`` wrapping the grid rather than a bare ``tuple`` subclass -- Pydantic + / ``Encodable`` need a real type with a core schema to move a ``State`` across the + model boundary (as a spliced prompt value and in the ``step`` signature). + + The grid itself is left unstructured -- recovering objects (player, box, walls) + from it is exactly the agent's job. ``grid`` is a tuple of tuples, so ``State`` is + ``frozen`` and hashable and doubles as a BFS key that compares by value. + """ + + grid: tuple[tuple[int, ...], ...] + + def __str__(self) -> str: + return "\n".join("".join(str(cell) for cell in row) for row in self.grid) + + +class Color(enum.IntEnum): + """The palette the agent observes. Only the *dynamics* are hidden; the goal -- + ``BOX_ON_TARGET`` appearing -- is visible, so we never synthesize an is_goal.""" + + FLOOR = 0 + WALL = 1 + PLAYER = 2 + BOX = 3 + TARGET = 4 + BOX_ON_TARGET = 5 + + +class Action(enum.IntEnum): + """The four moves. ``IntEnum`` so a value stays a plain ``int`` at the model + boundary: the synthesized ``step`` sees actions as 0-3, exactly as the prompt says.""" + + UP = 0 + DOWN = 1 + LEFT = 2 + RIGHT = 3 + + @property + def delta(self) -> tuple[int, int]: + return { + Action.UP: (-1, 0), + Action.DOWN: (1, 0), + Action.LEFT: (0, -1), + Action.RIGHT: (0, 1), + }[self] + + +class Transition(typing.NamedTuple): + """One recorded step of ground truth in the Timeline.""" + + before: State + action: Action + after: State + + +# A ``(row, column)`` cell coordinate in the grid. +type Position = tuple[int, int] + + +class Game(abc.ABC): + """A hidden game with a visible goal. The agent must reverse-engineer the rules.""" + + rows: int + cols: int + + @abc.abstractmethod + def observe(self) -> State: + """Return the current grid observation.""" + raise NotImplementedError + + @abc.abstractmethod + def step(self, action: Action) -> tuple[State, bool]: + """Apply an action to reality and return the new state and whether the goal is + reached.""" + raise NotImplementedError + + +class PushGame(Game): + """A tiny Sokoban-lite game. The player pushes a box onto a target. + + The true mechanics -- a move steps the player one cell; stepping into the box + pushes it one further; walls block both -- are *not* revealed to the agent. The + box can only travel rightward toward the target here, so the level has no dead + ends: the agent always recovers once its model is correct. + """ + + rows: int + cols: int + walls: set[Position] + player: Position + box: Position + targets: set[Position] + + def __init__(self) -> None: + self.rows, self.cols = 4, 7 + self.walls = { + (r, c) + for r in range(self.rows) + for c in range(self.cols) + if r in (0, self.rows - 1) or c in (0, self.cols - 1) + } + self.player = (1, 1) + self.box = (1, 2) + self.targets = {(1, 5)} + + def observe(self) -> State: + grid = [[Color.FLOOR] * self.cols for _ in range(self.rows)] + for r, c in self.walls: + grid[r][c] = Color.WALL + for r, c in self.targets: + grid[r][c] = Color.TARGET + br, bc = self.box + grid[br][bc] = Color.BOX_ON_TARGET if self.box in self.targets else Color.BOX + pr, pc = self.player + grid[pr][pc] = Color.PLAYER + return State(tuple(tuple(int(cell) for cell in row) for row in grid)) + + def step(self, action: Action) -> tuple[State, bool]: + dr, dc = action.delta + pr, pc = self.player + ahead = (pr + dr, pc + dc) + if ahead in self.walls: + pass # blocked by a wall + elif ahead == self.box: + beyond = (ahead[0] + dr, ahead[1] + dc) + if beyond not in self.walls: # push the box (never a wall here) + self.box = beyond + self.player = ahead + else: + self.player = ahead + return self.observe(), self.box in self.targets diff --git a/docs/source/llm_examples/reasoning/hanoi.py b/docs/source/llm_examples/reasoning/hanoi.py new file mode 100644 index 000000000..419e80c0e --- /dev/null +++ b/docs/source/llm_examples/reasoning/hanoi.py @@ -0,0 +1,227 @@ +"""LLM-based Towers of Hanoi solver with two strategies. + +Two solving strategies share a common ``Step`` / ``GameState`` model and are +selected with ``--mode``: + +- ``recursive`` — ask the LLM to return the full move list in one shot, using + the classic recursive decomposition. +- ``iterative`` — ask the LLM for one move at a time, with tool-based + validation. Adapted from https://github.com/BasisResearch/effectful/pull/404. + Demonstrates: + + - A static ``Step`` model for structured output + - ``@Tool.define`` inside a closure to expose game-state validation as a tool + - Templates defined inside a function that auto-capture closure-scoped tools +""" + +import argparse +import dataclasses +import itertools + +from effectful.handlers.llm import Template, Tool + + +@dataclasses.dataclass +class Step: + """A single move: take the top disk from tower ``start`` and place it on + tower ``end``. Tower indices are zero-based.""" + + start: int + end: int + explanation: str = dataclasses.field(default="") # optional reasoning from the LLM + + +@dataclasses.dataclass +class GameState: + """State of a Towers of Hanoi game. + + Higher numbers represent larger disks, so ``(2, 1, 0)`` is a valid + tower (largest on bottom). The goal is to move all disks from the + leftmost tower (index 0) to the rightmost tower (index -1). + + This is a plain ``dataclass`` (not a Pydantic model) so the type checker + can see its methods. + """ + + size: int + towers: tuple[tuple[int, ...], ...] = dataclasses.field(default=()) + + def __post_init__(self): + if self.size > 0 and not self.towers: + self.towers = tuple( + tuple(reversed(range(self.size))) if i == 0 else () + for i in range(self.size) + ) + + def apply(self, step: Step) -> "GameState": + """Apply a move, returning the new state. Raises ``ValueError`` if + the move is invalid.""" + start, end = step.start, step.end + if not (0 <= start < len(self.towers) and 0 <= end < len(self.towers)): + raise ValueError(f"tower index out of range: ({start}, {end})") + if len(self.towers[start]) == 0: + raise ValueError(f"tower {start} is empty") + if len(self.towers[end]) > 0 and self.towers[start][-1] > self.towers[end][-1]: + raise ValueError( + f"cannot place disk {self.towers[start][-1]} on top of " + f"disk {self.towers[end][-1]}" + ) + new_towers = [list(t) for t in self.towers] + disk = new_towers[start].pop() + new_towers[end].append(disk) + return GameState(self.size, tuple(tuple(t) for t in new_towers)) + + def is_done(self) -> bool: + return all(len(t) == 0 for t in self.towers[:-1]) and all( + self.towers[-1][i] > self.towers[-1][i + 1] + for i in range(len(self.towers[-1]) - 1) + ) + + def valid_steps(self) -> list[Step]: + steps = [] + for i, ti in enumerate(self.towers): + for j, tj in enumerate(self.towers): + if i == j or len(ti) == 0: + continue + if len(tj) == 0 or ti[-1] < tj[-1]: + steps.append(Step(i, j)) + return steps + + def __str__(self) -> str: + return " | ".join(str(list(t)) for t in self.towers) + + +# --------------------------------------------------------------------------- +# Recursive solver +# --------------------------------------------------------------------------- + + +def validate_solution(size: int, steps: list[Step]) -> bool: + """Apply all steps to the initial state and check that the puzzle is solved.""" + state = GameState(size=size) + print(f" initial: {state}") + for i, step in enumerate(steps): + try: + state = state.apply(step) + print(f" step {i}: move {step.start} -> {step.end} => {state}") + except ValueError as e: + print(f" step {i}: INVALID move {step.start} -> {step.end}: {e}") + return False + if state.is_done(): + print(f" Solved in {len(steps)} moves!") + return True + else: + print(f" Not solved after {len(steps)} moves. Final state: {state}") + return False + + +def solve_recursive(state: GameState) -> None: + + @Template.define + def solve(n_disks: int, source: int, target: int, auxiliary: int) -> list[Step]: + """Solve Tower of Hanoi using recursion: move {n_disks} disks from tower {source} to + tower {target}, using tower {auxiliary} as temporary storage. + """ + + size = state.size + print(f"Solving Tower of Hanoi with {size} disks...") + steps = solve(n_disks=size, source=0, target=size - 1, auxiliary=1) + print(f"\nLLM returned {len(steps)} steps. Validating...\n") + validate_solution(size, steps) + + +# --------------------------------------------------------------------------- +# Iterative solver +# --------------------------------------------------------------------------- + + +def predict_next_step(state: GameState) -> Step: + """Ask the LLM to predict the next move. + + A ``get_valid_moves`` tool is defined in the closure so the template + can query which moves are legal for the current game state. A + ``validate_move`` tool checks whether a proposed move is legal and + raises ``ValueError`` if not — when wrapped by ``RetryLLMHandler``, + this error is fed back to the LLM so it can correct itself. + """ + valid = state.valid_steps() + + @Tool.define + def get_valid_moves() -> list[Step]: + """Return the list of valid moves for the current game state.""" + return valid + + @Tool.define + def validate_move(proposed: Step) -> bool: + """Check whether moving from tower ``start`` to tower ``end`` is legal.""" + return proposed in state.valid_steps() + + @Template.define + def predict(game_state: GameState) -> Step: + """Given the state of the game of Towers of Hanoi: + + {game_state} + + Predict the next step to complete the game (move all disks to the + rightmost tower). You MUST call get_valid_moves first to see which + moves are legal, then pick the best one. Give a brief reasoning. + """ + + return predict(state) + + +def solve_iterative(state: GameState, *, max_steps: int = 30) -> None: + """Solve Towers of Hanoi by repeatedly asking the LLM for the next move.""" + for i in itertools.count(): + print(f"step {i}: {state}") + if state.is_done(): + print("Solved!") + return + if i >= max_steps: + print("Gave up after max steps.") + return + + step: Step = predict_next_step(state) + try: + state = state.apply(step) + print(f" move: {step.start} -> {step.end}") + except ValueError as e: + print(f" attempt {i}: invalid move {step}: {e}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--mode", + choices=("recursive", "iterative"), + default="iterative", + help="Solving strategy: full recursive solution or iterative one move at a time", + ) + parser.add_argument( + "--game-size", + type=int, + default=3, + help="Number of disks in the Towers of Hanoi game", + ) + parser.add_argument( + "--max-steps", + type=int, + default=30, + help="Maximum number of steps before giving up (iterative mode only)", + ) + args = parser.parse_args() + + state = GameState(size=args.game_size) + if args.mode == "recursive": + solve_recursive(state) + else: + solve_iterative(state, max_steps=args.max_steps) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/lineup.py b/docs/source/llm_examples/reasoning/lineup.py new file mode 100644 index 000000000..86b850234 --- /dev/null +++ b/docs/source/llm_examples/reasoning/lineup.py @@ -0,0 +1,108 @@ +"""Solving hard problems by writing and running Python. + +You are a careful problem solver and an expert Python programmer. You answer by +writing code, not by reasoning in prose alone: problems that are error-prone to +work out by hand are often easy to brute-force or verify with a short program. +""" + +import argparse +import collections.abc +import dataclasses +import typing + +from effectful.handlers.llm import Template + + +@dataclasses.dataclass(frozen=True) +class LineupClue: + """ + A clue about the relative ordering of n people, numbered 0 to n - 1, in a line. + Used to describe puzzles like the classic "zebra puzzle" represented in `solve_lineup`. + Each `LineupClue` corresponds to a single ordering constraint, ``(kind, a, b)``. + + The meaning of ``a`` and ``b`` depends on ``kind``: + + - ``("at", a, k)`` -- person ``a`` is at position ``k`` + - ``("left", a, b)`` -- person ``a`` is somewhere left of person ``b`` + - ``("imm_left", a, b)`` -- person ``a`` is immediately left of person ``b`` + - ``("adj", a, b)`` -- persons ``a`` and ``b`` are in adjacent positions + """ + + kind: typing.Literal["at", "left", "imm_left", "adj"] + a: int + b: int + + +@Template.define +def solve_lineup(n: int, clues: collections.abc.Sequence[LineupClue]) -> list[int]: + """Solve a 'zebra'-style ordering puzzle: place n={n} people, numbered 0 to + n - 1, in a line in positions 1 to n (each position used once) so that every + `LineupClue` in the following list holds: + + {clues} + + Every puzzle has exactly one consistent arrangement. Return the list of + positions ``[position of 0, position of 1, ..., position of n - 1]``, + as shown in the following worked examples: + + >>> solve_lineup(3, [LineupClue("at", 0, 1), LineupClue("left", 1, 2)]) + [1, 2, 3] + >>> solve_lineup(4, [LineupClue("left", 0, 1), LineupClue("left", 1, 2), LineupClue("left", 2, 3)]) + [1, 2, 3, 4] + >>> solve_lineup(4, [LineupClue("imm_left", 0, 1), LineupClue("at", 2, 4), LineupClue("left", 3, 0)]) + [2, 3, 4, 1] + >>> solve_lineup(5, [LineupClue("at", 0, 3), LineupClue("imm_left", 1, 2), LineupClue("left", 3, 4), LineupClue("at", 4, 5)]) + [3, 1, 2, 4, 5] + """ + + +def main() -> None: + kinds = typing.get_args(LineupClue.__annotations__["kind"]) + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--n", + type=int, + default=5, + help="Number of people in the line (used with --clue)", + ) + parser.add_argument( + "--clue", + dest="clues", + action="append", + nargs=3, + metavar=("KIND", "A", "B"), + default=None, + help=( + f"An ordering constraint 'KIND A B' where KIND is one of " + f"{'/'.join(kinds)} (e.g. --clue imm_left 0 1); repeatable" + ), + ) + args = parser.parse_args() + + if args.clues is not None: + n = args.n + clues = [] + for kind, a, b in args.clues: + if kind not in kinds: + parser.error( + f"invalid clue kind {kind!r}; choose from {'/'.join(kinds)}" + ) + try: + clues.append(LineupClue(kind, int(a), int(b))) + except ValueError: + parser.error(f"clue positions must be integers, got {a!r} {b!r}") + else: + n = 5 + clues = [ + LineupClue("imm_left", 0, 1), + LineupClue("imm_left", 1, 2), + LineupClue("at", 3, 5), + LineupClue("left", 4, 0), + ] + + print(f"Zebra-style ordering puzzle: n={n}, clues={clues}") + print(f"Answer: {solve_lineup(n, clues)}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/taboo.py b/docs/source/llm_examples/reasoning/taboo.py new file mode 100644 index 000000000..ffb7f02d0 --- /dev/null +++ b/docs/source/llm_examples/reasoning/taboo.py @@ -0,0 +1,168 @@ +"""Multi-agent Taboo word guessing game. + +Demonstrates: +- Two ``Agent`` instances with independent conversation histories +- Inter-agent communication via plain function calls +- Each agent has a different persona and goal +- ``Agent.__history__`` keeps each agent's context isolated +""" + +import argparse +import dataclasses +import enum + +from effectful.handlers.llm import Agent, Template, Tool + +# --------------------------------------------------------------------------- +# Structured output +# --------------------------------------------------------------------------- + + +class Confidence(enum.StrEnum): + LOW = "low" + MEDIUM = "medium" + HIGH = "high" + + +@dataclasses.dataclass(frozen=True) +class Guess: + guess: str + confidence: Confidence + + +# --------------------------------------------------------------------------- +# Agents +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Hinter(Agent): + """Agent that gives hints about a secret word without saying it.""" + + secret_word: str + taboo_words: list[str] + + @Tool.define + def is_taboo(self, hint: str) -> bool: + """Check if the given hint contains any taboo words or the secret word.""" + lowered_hint = hint.lower() + if self.secret_word.lower() in lowered_hint: + return True + for taboo in self.taboo_words: + if taboo.lower() in lowered_hint: + return True + return False + + @Template.define + def give_hint(self, guesser_response: str) -> str: + """You are playing a word guessing game. You must help the guesser + figure out the secret word by giving creative hints. + + RULES: + - You MUST NOT say the secret word: {self.secret_word} + - You MUST NOT use any of these taboo words: {self.taboo_words} + - Give a single, concise hint (one sentence) + - Review conversation history to avoid repeating hints + - Use the is_taboo tool to check if your hint is valid + + The guesser's last response was: {guesser_response} + """ + + +class Guesser(Agent): + """Agent that tries to guess the secret word from hints.""" + + @Template.define + def make_guess(self, hint: str) -> Guess: + """You are playing a word guessing game. Based on the hints you've + received, guess the secret word. + + Latest hint: {hint} + + Review the conversation history for all previous hints. + Make your best guess. + """ + + +# --------------------------------------------------------------------------- +# Game loop +# --------------------------------------------------------------------------- + + +def play_taboo( + secret_word: str, + taboo_words: list[str], + max_rounds: int = 5, +) -> bool: + """Play a round of Taboo between a hinter and a guesser.""" + hinter = Hinter(secret_word=secret_word, taboo_words=taboo_words) + guesser = Guesser() + + guesser_response = "I'm ready to guess!" + + for round_num in range(max_rounds): + # Hinter gives a hint + hint = hinter.give_hint(guesser_response) + print(f" [round {round_num}] Hinter: {hint}") + + # Guesser tries to guess + guess = guesser.make_guess(hint) + guesser_response = f"I guessed '{guess.guess}' ({guess.confidence})" + print(f" [round {round_num}] Guesser: {guess.guess} ({guess.confidence})") + + if guess.guess.lower().strip() == secret_word.lower(): + print(f" Correct! Guessed in {round_num} round(s).") + return True + + print(f" Failed to guess '{secret_word}' in {max_rounds} rounds.") + return False + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--max-rounds", + type=int, + default=5, + help="Maximum rounds per game", + ) + parser.add_argument( + "--secret-word", + type=str, + default=None, + metavar="WORD", + help="Secret word to guess (used with --taboo-words for a single custom game)", + ) + parser.add_argument( + "--taboo-words", + nargs="+", + type=str, + default=None, + metavar="WORD", + help="Taboo words the hinter may not say (used with --secret-word)", + ) + args = parser.parse_args() + + if (args.secret_word is None) != (args.taboo_words is None): + parser.error("--secret-word and --taboo-words must be given together") + + if args.secret_word is not None: + games = [(args.secret_word, args.taboo_words)] + else: + games = [ + ("piano", ["music", "keys", "instrument", "play"]), + ("volcano", ["lava", "eruption", "mountain", "hot"]), + ] + + for secret, taboo in games: + print(f"\nGame: '{secret}' (taboo: {taboo})") + play_taboo(secret, taboo, max_rounds=args.max_rounds) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/theory_of_mind.py b/docs/source/llm_examples/reasoning/theory_of_mind.py new file mode 100644 index 000000000..edde9db9f --- /dev/null +++ b/docs/source/llm_examples/reasoning/theory_of_mind.py @@ -0,0 +1,95 @@ +"""Solving hard problems by writing and running Python. + +You are a careful problem solver and an expert Python programmer. You answer by +writing code, not by reasoning in prose alone: problems that are error-prone to +work out by hand are often easy to brute-force or verify with a short program. +""" + +import argparse +import collections.abc + +from effectful.handlers.llm import Template + + +@Template.define +def musr_object_placement( + story: str, person: str, item: str, locations: collections.abc.Sequence[str] +) -> str: + """A MuSR object-placement question: a theory-of-mind puzzle. Read the story + and decide, from {locations}, where {person} would look for the {item}. + + The answer is the last place {person} *saw* the {item}: the last move they + watched, or any later moment they directly saw it somewhere; or its original + location if they never saw it after that. A person's belief does not change + while they are not watching, so where the {item} actually ends up and where + {person} believes it is can differ. + + {story} + + >>> musr_object_placement( + ... "Danny set the earphones in the recording booth, then stepped out for a " + ... "call. While he was gone, Emma quietly moved them to the producer's desk.", + ... "Danny", + ... "earphones", + ... ["recording booth", "producer's desk"], + ... ) + 'recording booth' + """ + + +def main() -> None: + STUDIO_STORY = """\ +In the heart of the bustling studio, Ricky, Emma, and Danny readied themselves \ +for a day of creating magic. Ricky, the gifted singer-songwriter, had his \ +precious notebook of lyrics on the producer's desk. Emma, their producer, was \ +cognizant of the notebook's place at her desk. Across the room, Danny, the studio \ +assistant, kept the earphones in the recording booth. They were all aware of the \ +arrangement -- the notebook on the producer's desk, the earphones in the \ +recording booth. + +Ricky gently places his notebook onto the piano, then becomes engrossed in \ +perfecting his song. Emma, engrossed in her thoughts, deftly moves the earphones \ +to the producer's desk. At that moment Danny was in a stirring conversation with a \ +visiting sound engineer; the visitor stood blocking Danny's general overview of \ +the studio space. + +Later, delicately lifting Ricky's notebook, Danny orchestrates its move to the \ +producer's desk. At the desk, he glimpses a pair of earphones indirectly drawing \ +his attention amidst his routine of tidying up. Meanwhile Emma, from inside a \ +sound-proofed booth, was lost in reviewing already-recorded tracks, out of \ +Danny's view.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--story", + type=str, + default=STUDIO_STORY, + help="The narrative describing where the item is moved and who saw it", + ) + parser.add_argument( + "--person", + type=str, + default="Danny", + help="The person whose belief about the item's location is queried", + ) + parser.add_argument( + "--item", + type=str, + default="earphones", + help="The object being tracked", + ) + parser.add_argument( + "--locations", + nargs="+", + default=["piano", "producer's desk", "recording booth"], + metavar="LOCATION", + help="Candidate locations to choose the answer from", + ) + args = parser.parse_args() + + answer = musr_object_placement(args.story, args.person, args.item, args.locations) + print(f"MuSR: where would {args.person} look for the {args.item}?") + print(f"Answer: {answer}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/reasoning/world_model_agent.py b/docs/source/llm_examples/reasoning/world_model_agent.py new file mode 100644 index 000000000..2d6886223 --- /dev/null +++ b/docs/source/llm_examples/reasoning/world_model_agent.py @@ -0,0 +1,227 @@ +"""Schema-style world-model agent: learn a hidden game by writing its rules as code. + +Inspired by the "Schema" harness (https://schema-harness.github.io/), which has a +model play a game with hidden rules "like a physicist": write the game's mechanism +as an executable program, test it against the recorded history, and plan inside it. + +Demonstrates: +- A ``Template`` returning a ``Callable`` -- the model's *world model* is executable + Python, synthesized once and then run thousands of times by ordinary code +- An ``Agent`` whose persistent state is the memory: an append-only Timeline of real + transitions that is fed back into every deliberation as ground truth +- Certification at synthesis: the model embeds recorded transitions as doctests in the + ``step`` it writes, and the ``Callable`` decoder runs them -- a model that fails to + reproduce recorded history is rejected and fed back before it is ever used +- Reality-outranks-model: during execution a single mispredict voids the rest of the + plan and appends the surprising transition, forcing a re-theorize next round; a + plain-Python BFS searches *inside* the synthesized model for free + +The Timeline is external memory fed into every deliberation as ground truth -- and as +the doctests that certify each model -- so it is genuinely distinct from the Agent's +conversational history. We deliberately keep +the *notes* store out of this (a) variant: because the tiny game never overflows the +context window, a rewritable notes summary would only duplicate ``__history__``. At +ARC-AGI-3 scale, where context is auto-compacted, a curated notes file stops being +redundant and becomes the model's "weights" -- that is the (b) variant this omits. +""" + +# Remaining simplifications vs. the source (beyond the (b)-variant notes store above): +# - State grounding is given, not discovered. Schema's headline claim is that the agent +# invents *which pixels are objects* (Level 1) jointly with the transition rule +# (Level 2) in one program. Here the palette is already semantically labelled (see +# ``gridworlds.Color``), so the agent only maps integers to roles and infers dynamics +# -- closer to WorldCoder's "rule over a given state" than to Schema's joint problem. +# - The goal predicate is hardcoded, not inferred. Schema synthesizes ``is_goal`` too; +# here ``plan`` tests for ``BOX_ON_TARGET`` directly, which is fair only because this +# game renders its goal as a visible color. +# - Exploration is a heuristic, not a discriminating experiment. Schema keeps several +# candidate rules and probes the action where they *predict different outcomes*; +# ``explore`` instead asks the model for one "informative" action, with no ensemble +# to disagree. +# - Certification is self-reported. Schema's ``run_backtest`` replays a model over the +# *entire* recorded history externally; here it rests on the model faithfully +# transcribing the "salient" transitions as doctests, which the decoder then runs. + +import argparse +import collections +import collections.abc +import dataclasses +import textwrap + +from gridworlds import Action, Color, Game, State, Transition + +from effectful.handlers.llm import Agent, Template + + +@dataclasses.dataclass +class Physicist(Agent): + """Reverse-engineers the game by writing its ``step`` rule as Python code.""" + + hint: str + timeline: list[Transition] = dataclasses.field(default_factory=list) + + @Template.define + def explore(self, state: State) -> Action: + """ + Propose an action that would be informative about the hidden dynamics, + given the current world state: + + + {state} + + + and the recorded transitions so far: + + + {self.timeline} + + + and the high-level hint about the game: + + + {self.hint} + + + Do not use any tools. + """ + + @Template.define + def theorize( + self, state: State + ) -> collections.abc.Callable[[State, Action], State]: + """You are reverse-engineering a 2D grid game by writing its rules as code. + You've been given a high-level hint about the game: + + + {self.hint} + + + Beyond that, the dynamics are hidden; infer them ONLY from these recorded transitions: + + + {self.timeline} + + + The current world state, which you will plan beyond using the model, is: + + + {state} + + + Write a pure function ``step(state, action)`` that reproduces every recorded transition exactly. + The function's docstring **MUST** include all salient recorded transitions + from the timeline as runnable doctests. If there are no recorded transitions, + you do not need to include any doctests. + """ + + def plan( + self, + model: collections.abc.Callable[[State, Action], State], + start: State, + *, + max_nodes: int = 5000, + ) -> list[Action]: + """Search *inside* the model for a plan reaching the goal (BOX_ON_TARGET). Free. + + Returns the action sequence to a goal state (``[]`` if ``start`` already wins), + or ``[]`` if no plan is found within ``max_nodes``. + """ + solved = lambda s: any(Color.BOX_ON_TARGET in row for row in s.grid) # noqa: E731 + if solved(start): + return [] + frontier: collections.deque[tuple[State, list[Action]]] = collections.deque( + [(start, [])] + ) + seen: set[State] = {start} + while frontier and len(seen) < max_nodes: + state, plan = frontier.popleft() + for action in Action: + try: + nxt = model(state, action) + except Exception: + continue # can't plan through a rule that crashes + if nxt in seen: + continue + if solved(nxt): + return plan + [action] + seen.add(nxt) + frontier.append((nxt, plan + [action])) + return [] + + def solve(self, env: Game, *, max_actions: int = 40) -> bool: + """ + Outer loop: observe, deliberate, plan, execute. + """ + while len(self.timeline) < max_actions: + # Observe the current state of reality and print it. + state = env.observe() + print(f"\ncurrent grid ({len(self.timeline)} real actions spent):\n{state}") + + # Deliberate: synthesize a step() model; its embedded doctests certify it + # against the recorded Timeline at decode time. + model = self.theorize(state) + + # Plan inside the certified model for free; if none, take one probing step. + plan = self.plan(model, state) + if not plan: + plan = [self.explore(state)] + print( + f"[plan] no solution in model; probing with action {plan[0].name}" + ) + else: + print(f"[plan] found in model: {[a.name for a in plan]}") + + # Execute against reality, checking each prediction. A surprise voids the rest. + for action in plan: + predicted = model(state, action) + actual, done = env.step(action) + self.timeline.append(Transition(state, action, actual)) + state = actual + if done: + print( + f"[execute] action {action.name} -> SOLVED in {len(self.timeline)} actions" + ) + return True + if actual != predicted: + print(f"[execute] action {action.name} -> surprise; plan voided") + break + print(f"[execute] action {action.name} -> as predicted") + + print(f"\nGave up after {len(self.timeline)} actions.") + return False + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--env", + type=str, + choices=[ + "push", + ], + default="push", + help="Which hidden game to solve", + ) + parser.add_argument( + "--max-actions", + type=int, + default=40, + help="Budget of real environment actions before giving up", + ) + args = parser.parse_args() + + if args.env == "push": + from docs.source.llm_examples.reasoning.gridworlds import PushGame + + game = PushGame() + else: + raise ValueError(f"Unknown environment {args.env}") + + assert game.__doc__, "Game must have a docstring hint for the agent" + phys = Physicist(hint=textwrap.dedent(game.__doc__)) + solved = phys.solve(game, max_actions=args.max_actions) + assert solved, "Failed to solve the game within the action budget." + + +if __name__ == "__main__": + main() diff --git a/effectful/handlers/llm/__init__.py b/effectful/handlers/llm/__init__.py index cdda93479..a8c107c4c 100644 --- a/effectful/handlers/llm/__init__.py +++ b/effectful/handlers/llm/__init__.py @@ -1,3 +1,43 @@ -from .template import Agent, Template, Tool +"""LLM-implemented functions via algebraic effects. -__all__ = ["Agent", "Template", "Tool"] +`effectful.handlers.llm` lets you write Python functions whose bodies are +implemented by a large language model, and call them like ordinary code. + +## Core concepts + +- **`Template`** — a fully type-annotated Python function whose body is `raise + NotHandled` and whose docstring is a [format + string](https://docs.python.org/3/library/string.html#format-string-syntax) + prompt. Calling a template (under a provider) formats its arguments into the + prompt, invokes the model, and decodes the response to the template's declared + return type. Define one with the `Template.define` decorator. + +- **`Tool`** — a normal Python callable exposed to the model. Its signature and + docstring become the schema the model sees; the model calls it by name with + JSON arguments and receives the encoded result. Tools in a template's lexical + scope are offered to the model automatically; because scope is ordinary Python + scope, an `Agent` (or an enclosing function) naturally partitions tools and + templates into disjoint sets. Define one with `Tool.define`. + +- **`Agent`** — a class mixin giving each instance a persistent message history, + so its `Template` methods accumulate conversation context across calls. + Instance attributes are available in prompts via `{self.attr}`. + +- **`Encodable`** — the type-driven JSON bridge used internally to encode Python + values into the model's context and decode the model's output (structured + return values and tool-call arguments) back into typed Python objects. + +## Tool calling and structured output + +During a template call the model may take multiple turns: on each turn it can +call any `Tool` in scope (results are fed back and the loop continues) or +produce a final answer. The final answer is decoded to the template's return +type via constrained/structured generation, so non-`str` return types (ints, +dataclasses, etc.) come back as real Python values. A `FinalTool` lets the model +"answer" by calling a tool whose return value becomes the result and terminates +the loop. +""" + +from .types import Agent, Encodable, Template, Tool + +__all__ = ["Agent", "Template", "Tool", "Encodable"] diff --git a/effectful/handlers/llm/completions.py b/effectful/handlers/llm/completions.py deleted file mode 100644 index a393c7d90..000000000 --- a/effectful/handlers/llm/completions.py +++ /dev/null @@ -1,751 +0,0 @@ -import abc -import collections -import collections.abc -import dataclasses -import functools -import inspect -import json -import string -import textwrap -import traceback -import typing -import uuid - -import litellm -import pydantic -import tenacity -from litellm import ( - ChatCompletionFunctionMessage, - ChatCompletionMessageToolCall, - ChatCompletionTextObject, - ChatCompletionToolMessage, - OpenAIChatCompletionAssistantMessage, - OpenAIChatCompletionSystemMessage, - OpenAIChatCompletionUserMessage, - OpenAIMessageContentListBlock, -) - -from effectful.handlers.llm.encoding import ( - REPL_ANCHOR_KEY, - TYPE_CHECK_ANCHOR_KEY, - DecodedToolCall, - Encodable, - to_content_blocks, -) -from effectful.handlers.llm.evaluation import ReplSession, _repl_session -from effectful.handlers.llm.template import ( - Agent, - Template, - Tool, - _is_recursive_signature, -) -from effectful.internals.unification import nested_type -from effectful.ops.semantics import fwd, handler -from effectful.ops.syntax import ObjectInterpretation, implements -from effectful.ops.types import Operation - - -class AssistantMessage(OpenAIChatCompletionAssistantMessage): - id: str - - -class ToolMessage(ChatCompletionToolMessage): - id: str - - -class FunctionMessage(ChatCompletionFunctionMessage): - id: str - - -class SystemMessage(OpenAIChatCompletionSystemMessage): - id: str - - -class UserMessage(OpenAIChatCompletionUserMessage): - id: str - - -Message = AssistantMessage | ToolMessage | FunctionMessage | SystemMessage | UserMessage - -DEFAULT_SYSTEM_PROMPT = ( - "You are a helpful assistant, you need to follow user's instruction" -) - - -class _NoActiveHistoryException(Exception): - """Raised when there is no active message history to append to.""" - - -@Operation.define -def _get_history() -> collections.OrderedDict[str, Message]: - raise _NoActiveHistoryException( - "No active message history. This operation should only be used within a handler that provides a message history." - ) - - -def append_message(message: Message, last: bool = True) -> None: - try: - _get_history()[message["id"]] = message - if not last: - _get_history().move_to_end(message["id"], last=False) - except _NoActiveHistoryException: - pass - - -def _make_message(content: dict) -> Message: - m_id = content.get("id") or str(uuid.uuid1()) - message = typing.cast(Message, {**content, "id": m_id}) - return message - - -class DecodingError[E: Exception](abc.ABC, Exception): - """Base class for decoding errors that can occur during LLM response processing.""" - - original_error: E - - @abc.abstractmethod - def to_feedback_message(self, include_traceback: bool) -> Message: - """Convert the decoding error into a feedback message to be sent back to the LLM.""" - raise NotImplementedError - - -@dataclasses.dataclass -class ToolCallDecodingError[E: Exception](DecodingError[E]): - """Error raised when decoding a tool call fails.""" - - original_error: E - raw_message: Message - raw_tool_call: ChatCompletionMessageToolCall - - def __str__(self) -> str: - return f"Error decoding tool call '{self.raw_tool_call.function.name}': {self.original_error}. Please provide a valid response and try again." - - def to_feedback_message(self, include_traceback: bool) -> Message: - error_message = f"{self}" - if include_traceback: - tb = traceback.format_exc() - error_message = f"{error_message}\n\nTraceback:\n```\n{tb}```" - return _make_message( - { - "role": "tool", - "tool_call_id": self.raw_tool_call.id, - "content": error_message, - }, - ) - - -@dataclasses.dataclass -class ResultDecodingError[E: Exception](DecodingError[E]): - """Error raised when decoding the LLM response result fails.""" - - original_error: E - raw_message: Message - - def __str__(self) -> str: - return f"Error decoding response: {self.original_error}. Please provide a valid response and try again." - - def to_feedback_message(self, include_traceback: bool) -> Message: - error_message = f"{self}" - if include_traceback: - tb = traceback.format_exc() - error_message = f"{error_message}\n\nTraceback:\n```\n{tb}```" - return _make_message( - {"role": "user", "content": error_message}, - ) - - -@dataclasses.dataclass -class ToolCallExecutionError[E: Exception, T](DecodingError[E]): - """Error raised when a tool execution fails at runtime.""" - - original_error: E - raw_tool_call: DecodedToolCall[T] - - def __str__(self) -> str: - return f"Tool execution failed: Error executing tool '{self.raw_tool_call.name}': {self.original_error}" - - def to_feedback_message(self, include_traceback: bool) -> Message: - error_message = f"{self}" - if include_traceback: - tb = traceback.format_exc() - error_message = f"{error_message}\n\nTraceback:\n```\n{tb}```" - return _make_message( - { - "role": "tool", - "tool_call_id": self.raw_tool_call.id, - "content": error_message, - }, - ) - - -type MessageResult[T] = tuple[Message, typing.Sequence[DecodedToolCall], T | None] - -CACHE_CONTROL_EPHEMERAL = {"type": "ephemeral"} - - -def _add_cache_control_to_history( - history: collections.OrderedDict[str, "Message"], -) -> None: - """Add cache_control to the last user/tool message in an agent's history. - - This enables prompt caching on providers that support it (e.g. Anthropic). - Providers that don't support it (e.g. OpenAI) have cache_control stripped - by litellm's request transformation, so this is always safe to apply. - - Mutates the history OrderedDict in place. - """ - if not history: - return - for key in history: - msg = history[key] - if msg["role"] not in ("user", "tool", "assistant"): - continue - content = msg.get("content") - if isinstance(content, list) and content: - last_block = content[-1] - if isinstance(last_block, dict) and "cache_control" not in last_block: - new_content = list(content) - new_content[-1] = { - **last_block, - "cache_control": CACHE_CONTROL_EPHEMERAL, - } - history[key] = typing.cast(Message, {**msg, "content": new_content}) - - -class _LexicalVariableTool[T](Tool[[], T]): - """A zero-arg `Tool` that returns the captured value of a variable - from a `Template`'s lexical context. - - Tools are constructed fresh each `call_assistant` invocation, so - the reader closes over the snapshot `value` rather than the - surrounding `env` — in-place mutation of a mutable value is still - visible (same object reference), but rebinding the source name is - not. - """ - - @classmethod - def define(cls, value: typing.Any, *, name: str) -> "Tool[[], typing.Any]": - """Construct a synthetic reader Tool that returns `value`. - - Raises if `Encodable[nested_type(value)]` cannot be generated. - The caller is responsible for catching the failure and deciding - whether to skip the symbol. - """ - assert not isinstance(value, Tool), ( - "Tools are real tools and must not be re-wrapped as lexical readers." - ) - typ: typing.Any = nested_type(value).value - # Probe schema generation; raises if `Encodable[typ]` is not implemented. - pydantic.TypeAdapter(Encodable[typ]).json_schema() - - def tool_fn(): - return value - - tool_fn.__name__ = name - tool_fn.__qualname__ = name - tool_fn.__module__ = type(value).__module__ - tool_fn.__doc__ = ( - f"Reads the value of lexical variable `{name}` from the " - f"enclosing scope where this Template was defined. Takes " - f"no arguments; returns the current value." - ) - tool_fn.__annotations__ = {"return": typ} - return super().define(tool_fn) - - -@Operation.define -def collect_tools( - env: collections.abc.Mapping[str, typing.Any], -) -> collections.abc.Mapping[str, Tool]: - """Return the tools available to a Template given its lexical context. - - Default rule: real `Tool` and `Template` values bound directly in - `env`, plus `Tool` methods discovered through the MRO of any - `Agent` instance in `env`. Same-Tool-under-different-names is - deduped so each Tool appears exactly once. - - Handlers (see :class:`LexicalReaders`) may override this to add - synthetic readers, hide tools, etc. - """ - result: dict[str, Tool] = {} - - for name, obj in env.items(): - if isinstance(obj, Tool | Template): - result[name] = obj - elif isinstance(obj, Agent): - for cls in type(obj).__mro__: - for attr_name in vars(cls): - if isinstance(getattr(obj, attr_name), Tool): - result[f"{name}__{attr_name}"] = getattr(obj, attr_name) - - # Same Tool can appear under multiple names when visible both in the - # enclosing scope and via an Agent instance's MRO. Keep only the - # last name for each unique tool object. - tool2name = {tool: name for name, tool in sorted(result.items())} - for name, tool in tuple(result.items()): - if tool2name[tool] != name: - del result[name] - - return result - - -class LexicalReaders(ObjectInterpretation): - """Override `collect_tools` to also expose plain values from the - lexical context as zero-argument read-only Tools. Each non-Tool, - non-Template, non-Agent value bound to a valid identifier is - wrapped via `_LexicalVariableTool` if `Encodable[T]` accepts it; - schema-generation failures cause the symbol to be skipped. - """ - - @implements(collect_tools) - def _collect( - self, env: collections.abc.Mapping[str, typing.Any] - ) -> collections.abc.Mapping[str, Tool]: - result = dict(fwd()) - for name, obj in env.items(): - if name in result or not name.isidentifier(): - continue - try: - result[name] = _LexicalVariableTool.define(obj, name=name) - # `TypeError` joins the three Pydantic errors because the - # `Encodable[T]` registry raises `TypeError` to signal - # "no schema possible" — e.g. `_pydantic_type_operation`, - # `_pydantic_type_term`, and `_pydantic_callable`'s - # incomplete-signature path. Same intent as the Pydantic - # cases, different exception class. - except ( - pydantic.errors.PydanticSchemaGenerationError, - pydantic.errors.PydanticInvalidForJsonSchema, - pydantic.errors.PydanticUserError, - TypeError, - ): - continue - return result - - -class PythonRepl(ObjectInterpretation): - """Expose a persistent Python session to the LLM as an `exec_code` Tool. - - Off by default; install it where the LLM should be able to run code whose - state (variables, imports, definitions) survives across tool calls within a - single Template invocation. - - Scoping mirrors how `__history__` is managed for Template calls: `PythonRepl` - handles `Template.__apply__` to introduce a fresh `_repl_session` handler for - the duration of the call, and handles `collect_tools` to inject an `exec_code` - Tool routed to that session. The session is therefore introduced and - eliminated by its own handler, bounded to the Template call by construction -- - there is no global registry of sessions, and nested Template calls get their - own isolated sessions. - - The session is seeded from the Template's lexical context and routes execution - through the `parse`/`compile`/`exec` effect operations, so it works under any - installed eval provider (`UnsafeEvalProvider` or `RestrictedEvalProvider`). - """ - - @implements(Template.__apply__) - def _apply[**P, T]( - self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs - ) -> T: - # One session per Template call, created lazily on first use (the call's - # `env`, supplied by `collect_tools`/`exec_code`, seeds it). The - # enclosing `handler(...)` bounds the session's lifetime to this call, so - # nested Template calls introduce their own fresh session. - session: ReplSession | None = None - - def session_for( - env: collections.abc.MutableMapping[str, typing.Any], - ) -> ReplSession: - nonlocal session - if session is None: - session = ReplSession(env) - return session - - with handler({_repl_session: session_for}): - return fwd() - - @implements(collect_tools) - def _collect( - self, env: collections.abc.Mapping[str, typing.Any] - ) -> collections.abc.Mapping[str, Tool]: - tools = dict(fwd()) - # `collect_tools` only promises a `Mapping`, but the per-call `env` is the - # writable `ChainMap` the session splices its shared scope layer into, so - # narrow it for `_repl_session`/`ReplSession`. - tools["exec_code"] = _repl_session( - typing.cast(collections.abc.MutableMapping[str, typing.Any], env) - ).exec_code - return tools - - -@Operation.define -@functools.wraps(litellm.completion) -def completion(*args, **kwargs) -> typing.Any: - """Low-level LLM request. Handlers may log/modify requests and delegate via fwd(). - - This effect is emitted for model request/response rounds so handlers can - observe/log requests. - - """ - return litellm.completion(*args, **kwargs) - - -class _BoxedResponse[T](pydantic.BaseModel): - value: T - - -@Operation.define -def call_assistant[T]( - env: collections.abc.Mapping[str, typing.Any], - response_type: type[T], - model: str, - **kwargs, -) -> MessageResult[T]: - """Low-level LLM request. Handlers may log/modify requests and delegate via fwd(). - - This effect is emitted for model request/response rounds so handlers can - observe/log requests. - - Raises: - ToolCallDecodingError: If a tool call cannot be decoded. The error - includes the raw assistant message for retry handling. - ResultDecodingError: If the result cannot be decoded. The error - includes the raw assistant message for retry handling. - """ - anchor = kwargs.pop("anchor", None) # ride in kwargs; pop before the LLM call - tools = dict(collect_tools(env)) - tool_specs = { - k: typing.cast( - pydantic.TypeAdapter[typing.Any], - pydantic.TypeAdapter(Encodable[type(t)]), # type: ignore[misc] - ).dump_python(t, mode="json", context={k: t}) - for k, t in tools.items() - } - - # The OpenAI API requires a wrapper object for non-object structured output types, - # so we create one on the fly here. Using a Pydantic model offloads JSON schema - # generation and validation logic to litellm, and offers better error messages. - response_format: type[_BoxedResponse[T]] = pydantic.create_model( - "BoxedResponse", - value=Encodable[response_type], # type: ignore[valid-type] - __base__=_BoxedResponse, - ) - - response: litellm.types.utils.ModelResponse = completion( - model, - messages=list(_get_history().values()), - response_format=None if response_type is str else response_format, - tools=list(tool_specs.values()), - **kwargs, - ) - choice = response.choices[0] - assert isinstance(choice, litellm.types.utils.Choices) - - message: litellm.Message = choice.message - assert message.role == "assistant" - - raw_message = _make_message({**message.model_dump(mode="json")}) - append_message(raw_message) - - tool_calls: list[DecodedToolCall] = [] - encoding: pydantic.TypeAdapter[DecodedToolCall] = pydantic.TypeAdapter( - Encodable[DecodedToolCall] - ) - # Thread the type-check anchor into the tool-argument context under REPL_ANCHOR_KEY, so - # the `Encodable[CodeType]` decoder type-checks a `code` argument (the REPL `exec_code` - # tool) against the Template body at decode, splicing in the accumulated REPL session. - tool_context = {**tools, REPL_ANCHOR_KEY: anchor} if anchor is not None else tools - for raw_tool_call in message.get("tool_calls") or []: - try: - tool_calls += [ - encoding.validate_python(raw_tool_call, context=tool_context) - ] - except Exception as e: - raise ToolCallDecodingError( - raw_tool_call=raw_tool_call, - original_error=e, - raw_message=raw_message, - ) from e - - result = None - if not tool_calls: - # return response - serialized_result = message.get("content") or message.get("reasoning_content") - assert isinstance(serialized_result, str), ( - "final response from the model should be a string" - ) - if response_type is str: - result = typing.cast(T, serialized_result) - else: - try: - # Add the type-check anchor to the decode context only (not `env`, - # which is exposed as tools), so a synthesized result is checked - # against the Template's source. - result = response_format.model_validate( - json.loads(serialized_result), - context={**env, TYPE_CHECK_ANCHOR_KEY: anchor}, - ).value - except Exception as e: - raise ResultDecodingError(e, raw_message=raw_message) from e - - return (raw_message, tool_calls, result) - - -@Operation.define -def call_tool(tool_call: DecodedToolCall) -> Message: - """Implements a roundtrip call to a python function. Input is a json - string representing an LLM tool call request parameters. The output is - the serialised response to the model. - - """ - # call tool with python types - try: - result = tool_call.tool( - *tool_call.bound_args.args, **tool_call.bound_args.kwargs - ) - except Exception as e: - raise ToolCallExecutionError(raw_tool_call=tool_call, original_error=e) from e - - return_type: pydantic.TypeAdapter[typing.Any] = pydantic.TypeAdapter( - Encodable[nested_type(result).value] # type: ignore[misc] - ) - encoded_result = to_content_blocks( - return_type.dump_python(result, mode="json", context={}) - ) - message = _make_message( - dict(role="tool", content=encoded_result, tool_call_id=tool_call.id), - ) - append_message(message) - return message - - -@Operation.define -def call_user( - template: str, - env: collections.abc.Mapping[str, typing.Any], -) -> Message: - """ - Format a template applied to arguments into a user message. - """ - formatter = string.Formatter() - parts: list[OpenAIMessageContentListBlock] = [] - - buf: list[str] = [] - - def flush_text() -> None: - if buf: - parts.append(ChatCompletionTextObject(type="text", text="".join(buf))) - buf.clear() - - for literal, field_name, format_spec, conversion in formatter.parse( - textwrap.dedent(template) - ): - if literal: - buf.append(literal) - - if field_name is None: - continue - - obj, _ = formatter.get_field(field_name, (), env) - encoder: pydantic.TypeAdapter[typing.Any] = pydantic.TypeAdapter( - Encodable[nested_type(obj).value] # type: ignore[misc] - ) - encoded_obj = encoder.dump_python(obj, mode="json", context=env) - for part in to_content_blocks(encoded_obj): - if part["type"] == "text": - text = ( - formatter.convert_field(part["text"], conversion) - if conversion - else part["text"] - ) - buf.append(formatter.format_field(text, format_spec or "")) - else: - flush_text() - parts.append(part) - - flush_text() - - # Note: The OpenAI api only seems to accept images in the 'user' role. The - # effect of different roles on the model's response is currently unclear. - message = _make_message(dict(role="user", content=parts)) - append_message(message) - return message - - -@Operation.define -def call_system(template: Template) -> Message: - """Get system instruction message(s) to prepend to all LLM prompts.""" - system_prompt = template.__system_prompt__ or DEFAULT_SYSTEM_PROMPT - message = _make_message( - dict( - role="system", - content=[ - { - "type": "text", - "text": system_prompt, - "cache_control": {"type": "ephemeral"}, - } - ], - ) - ) - append_message(message, last=False) - return message - - -class RetryLLMHandler(ObjectInterpretation): - """Retries LLM requests if tool call or result decoding fails. - - This handler intercepts `call_assistant` and catches `ToolCallDecodingError` - and `ResultDecodingError`. When these errors occur, it appends error feedback - to the messages and retries the request. Malformed messages from retry attempts - are pruned from the final result. - - For runtime tool execution failures (handled via `call_tool`), errors are - captured and returned as tool response messages. - - Args: - include_traceback: If True, include full traceback in error feedback - for better debugging context (default: True). - catch_tool_errors: Exception type(s) to catch during tool execution. - Can be a single exception class or a tuple of exception classes. - Defaults to Exception (catches all exceptions). - stop: tenacity stop condition for retrying `call_assistant`. Defaults to - `tenacity.stop_after_attempt(4)`, which stops after 4 attempts. - **kwargs: Additional keyword arguments forwarded to `tenacity.Retrying`. - """ - - call_assistant_retryer: tenacity.Retrying - - _user_before_sleep: collections.abc.Callable[[tenacity.RetryCallState], None] | None - - def __init__( - self, - include_traceback: bool = True, - catch_tool_errors: type[BaseException] - | tuple[type[BaseException], ...] = Exception, - stop: tenacity.stop.stop_base = tenacity.stop_after_attempt(4), - **kwargs, - ): - self.include_traceback = include_traceback - self.catch_tool_errors = catch_tool_errors - assert "retry" not in kwargs, "Cannot override retry logic of RetryLLMHandler" - assert "reraise" not in kwargs, ( - "Cannot override reraise logic of RetryLLMHandler" - ) - self._user_before_sleep = kwargs.pop("before_sleep", None) - self.call_assistant_retryer = tenacity.Retrying( - retry=tenacity.retry_if_exception_type( - (ToolCallDecodingError, ResultDecodingError) - ), - reraise=True, - before_sleep=self._before_sleep, - stop=stop, - **kwargs, - ) - - def _before_sleep(self, retry_state: tenacity.RetryCallState) -> None: - e = retry_state.outcome.exception() # type: ignore - assert isinstance(e, (ToolCallDecodingError, ResultDecodingError)) - append_message(e.raw_message) - append_message(e.to_feedback_message(self.include_traceback)) - if self._user_before_sleep is not None: - self._user_before_sleep(retry_state) - - @implements(call_assistant) - def _call_assistant[T]( - self, - env: collections.abc.Mapping[str, typing.Any], - response_type: type[T], - model: str, - **kwargs, - ) -> MessageResult[T]: - _message_sequence = _get_history().copy() - - with handler({_get_history: lambda: _message_sequence}): - message, tool_calls, result = self.call_assistant_retryer(fwd) - - append_message(message) - return (message, tool_calls, result) - - @implements(call_tool) - def _call_tool(self, tool_call: DecodedToolCall) -> Message: - """Handle tool execution with runtime error capture. - - Runtime errors from tool execution are captured and returned as - error messages to the LLM. Only exceptions matching `catch_tool_errors` - are caught; others propagate up. - """ - try: - return fwd(tool_call) - except ToolCallExecutionError as e: - if isinstance(e.original_error, self.catch_tool_errors): - message = e.to_feedback_message(self.include_traceback) - append_message(message) - return message - else: - raise - - -class LiteLLMProvider(ObjectInterpretation): - """Implements templates using the LiteLLM API.""" - - config: collections.abc.Mapping[str, typing.Any] - - def __init__(self, model="gpt-4o", **config): - self.config = { - "model": model, - **inspect.signature(litellm.completion).bind_partial(**config).kwargs, - } - - @implements(Template.__apply__) - def _call[**P, T]( - self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs - ) -> T: - # encode arguments - bound_args = inspect.signature(template).bind(*args, **kwargs) - bound_args.apply_defaults() - env = template.__context__.new_child(bound_args.arguments) - - if not _is_recursive_signature(template.__signature__): - env = env.new_child({k: None for k, v in env.items() if v is template}) - - history: collections.OrderedDict[str, Message] = getattr( - template, "__history__", collections.OrderedDict() - ) # type: ignore - is_agent = hasattr(template, "__history__") - history_copy = history.copy() - - with handler({_get_history: lambda: history_copy}): - if ( - not _get_history() - or next(iter(_get_history().values()))["role"] != "system" - ): - call_system(template) - - message: Message = call_user(template.__prompt_template__, env) - - # For agents with persistent history, add cache_control to the - # last user message so the growing prefix gets cached on providers - # that support it (Anthropic). litellm strips it for OpenAI. - if is_agent: - _add_cache_control_to_history(history_copy) - - # loop based on: https://cookbook.openai.com/examples/reasoning_function_calls - tool_calls: list[DecodedToolCall] = [] - result: T | None = None - while message["role"] != "assistant" or tool_calls: - message, tool_calls, result = call_assistant( - env, - template.__signature__.return_annotation, - anchor=template.__default__, - **self.config, - ) - for tool_call in tool_calls: - message = call_tool(tool_call) - - try: - _get_history() - except _NoActiveHistoryException: - history.clear() - history.update(history_copy) - return typing.cast(T, result) diff --git a/effectful/handlers/llm/encoding.py b/effectful/handlers/llm/encoding.py deleted file mode 100644 index 4ccc8fd90..000000000 --- a/effectful/handlers/llm/encoding.py +++ /dev/null @@ -1,803 +0,0 @@ -import ast -import base64 -import dataclasses -import functools -import inspect -import io -import json -import linecache -import textwrap -import types -import typing -import uuid -from collections.abc import ( - Callable, - Mapping, - MutableMapping, -) -from typing import Any - -import litellm -import pydantic -from litellm import ( - ChatCompletionImageObject, - ChatCompletionMessageToolCall, - ChatCompletionTextObject, - ChatCompletionToolParam, - OpenAIMessageContentListBlock, -) -from openai.lib._pydantic import _ensure_strict_json_schema -from openai.types.chat import ( - ChatCompletionMessageToolCall as OpenAIChatCompletionMessageToolCall, -) -from PIL import Image - -import effectful.handlers.llm.evaluation as evaluation -from effectful.handlers.llm.template import Tool -from effectful.internals.unification import GenericAlias, TypeEvaluator, nested_type -from effectful.ops.types import Operation, Term - -type ToolCallID = str - -# Reserved key under which the type-check anchor (the enclosing Template's -# underlying function) rides in the Pydantic decoding context, alongside the -# lexical environment. `decode` reads it to type-check a synthesized function -# against the Template's source; absent (tool-argument decoding) means skip. -# Deliberately not a valid identifier so `LexicalReaders` skips it (no tool leak) -# and it can never collide with a lexical name. -TYPE_CHECK_ANCHOR_KEY = "" - -# Type-check anchor for REPL `exec_code` snippets, separate from the Callable/result -# synthesis anchor (TYPE_CHECK_ANCHOR_KEY): the two decoders check against different -# contracts -- a REPL snippet against the Template body, a synthesized Callable tool -# argument against its own parameter type. -REPL_ANCHOR_KEY = "" - -CONTENT_BLOCK_TYPES: frozenset[str] = frozenset( - literal - for member in typing.get_args(OpenAIMessageContentListBlock) - for literal in typing.get_args(typing.get_type_hints(member).get("type", str)) - if isinstance(literal, str) -) - - -@pydantic.validate_call(validate_return=True) -def to_content_blocks(value: typing.Any) -> list[OpenAIMessageContentListBlock]: - """Convert an encoded JSON-compatible value into a flat list of content blocks. - - Walks the value tree, extracting content-block-shaped dicts (identified by - their ``type`` discriminator) and emitting JSON syntax as text around them. - - Top-level strings are emitted bare (for natural template rendering). - Inside JSON structures, separators match ``json.dumps`` defaults so that - the linearization law holds for non-string encoded values: - ``linearize(to_content_blocks(v)) == json.dumps(v)``. - """ - if isinstance(value, str): - return [ChatCompletionTextObject(type="text", text=value)] - - buf: list[str] = [] - blocks: list[OpenAIMessageContentListBlock] = [] - - def flush() -> None: - if buf: - blocks.append(ChatCompletionTextObject(type="text", text="".join(buf))) - buf.clear() - - def walk(v: typing.Any) -> None: - if isinstance(v, dict) and v.get("type") in CONTENT_BLOCK_TYPES: - flush() - blocks.append(typing.cast(OpenAIMessageContentListBlock, v)) - elif isinstance(v, dict): - buf.append("{") - for i, (k, val) in enumerate(v.items()): - if i: - buf.append(", ") - buf.append(json.dumps(k) + ": ") - walk(val) - buf.append("}") - elif isinstance(v, list): - buf.append("[") - for i, item in enumerate(v): - if i: - buf.append(", ") - walk(item) - buf.append("]") - else: - buf.append(json.dumps(v)) - - walk(value) - flush() - return blocks - - -@dataclasses.dataclass(frozen=True, eq=True) -class DecodedToolCall[T]: - """ - Structured representation of a tool call decoded from an LLM response. - """ - - tool: Tool[..., T] - bound_args: inspect.BoundArguments - id: ToolCallID - name: str - - -if typing.TYPE_CHECKING: - type Encodable[T] = typing.Annotated[T, "encoded"] -else: - - class Encodable: - def __class_getitem__(cls, item): - return TypeToPydanticType().evaluate(item) - - -class TypeToPydanticType(TypeEvaluator): - """Substitute custom types with their Pydantic Annotated equivalents. - - Recursively walks a type annotation tree, replacing leaf types that have - registered Pydantic annotations (e.g., Image.Image -> PydanticImage) and - reconstructing the full generic type. - - The result can be passed to pydantic.TypeAdapter() for automatic - validation and serialization of nested structures. - """ - - @staticmethod - @functools.singledispatch - def _registry(ty: type): - raise RuntimeError("should not be here!") - - @classmethod - def register(cls, *args, **kwargs): - return cls._registry.register(*args, **kwargs) - - def evaluate(self, ty): - app = super().evaluate(ty) - origin = typing.get_origin(app) - # Only dispatch on regular types. Special forms (Literal, Annotated, - # Union) have non-type origins that singledispatch can't resolve; pass - # them through for Pydantic to handle natively. - if isinstance(app, type | GenericAlias) and ( - origin is None or isinstance(origin, type) - ): - return self._registry.dispatch(origin or app)(app) - else: - return app - - -@TypeToPydanticType.register(str) -def _pydantic_type_str[T](ty: type[T]) -> type[T]: - return ty - - -@TypeToPydanticType.register(object) -def _pydantic_type_base(ty: type) -> Any: - return ty - - -class _ComplexModel(typing.TypedDict): - real: float - imag: float - - -@pydantic.validate_call(validate_return=True) -def _validate_complex(value: _ComplexModel) -> complex: - return complex(value["real"], value["imag"]) - - -@pydantic.validate_call(validate_return=True) -def _serialize_complex(value: complex) -> _ComplexModel: - return {"real": value.real, "imag": value.imag} - - -@TypeToPydanticType.register(complex) -def _pydantic_type_complex(ty): - """Encode ``complex`` as ``{"real": float, "imag": float}``.""" - - adapted_schema = pydantic.TypeAdapter(_ComplexModel).json_schema() - - return typing.Annotated[ - ty, - pydantic.PlainValidator(_validate_complex), - pydantic.PlainSerializer(_serialize_complex), - pydantic.WithJsonSchema({**adapted_schema, "additionalProperties": False}), - ] - - -_CODE_FILENAME_PREFIX = " types.CodeType: - if isinstance(value, types.CodeType): - return value - if not isinstance(value, str): - raise ValueError( - f"expected Python source as a string, got {type(value).__name__}" - ) - filename = f"{_CODE_FILENAME_PREFIX}{uuid.uuid4()}>" - try: - module = evaluation.parse(value, filename) - # Reject `__future__`/star imports: both are `SyntaxError` once nested in a - # function body, so such a snippet can't be spliced into the Template for - # type checking. - evaluation.scan_non_nestable(module) - except (SyntaxError, ValueError) as exc: - raise ValueError(f"source is not valid REPL code: {exc}") from exc - - # Type-check the snippet in its execution context, exactly as a synthesized - # `Callable` is (see `_pydantic_callable`): when the enclosing Template is the - # type-check anchor in the decode context, splice the accumulated REPL session (the - # `_repl_session` op is in scope during the response decode) plus this snippet into - # the Template body and check it. A type error raises here -> the tool-call decode - # fails -> `RetryLLMHandler` retries, so ill-typed code never reaches `runcode`. - ctx = info.context or {} - anchor = ctx.get(REPL_ANCHOR_KEY) - if anchor is not None: - # Pass an empty env (not `ctx`): the managed session ignores it, and a fresh - # fallback session must not be seeded from the decode context (which holds tool - # names and the anchor key). The decoder only reads `prior_snippets`. - prior = evaluation._repl_session({}).prior_snippets - checked = evaluation._splice_repl(prior, value, anchor) - if checked is not None: - evaluation.type_check(*checked, lenient=True) - try: - return evaluation.compile(module, filename) - except (SyntaxError, ValueError) as exc: - raise ValueError(f"source does not compile: {exc}") from exc - - return typing.Annotated[ - ty, - pydantic.PlainValidator(validate), - pydantic.PlainSerializer( - lambda value: "".join(linecache.getlines(value.co_filename)) - ), - pydantic.WithJsonSchema({"type": "string"}), - ] - - -def _inline_refs(schema: dict) -> dict: - """Inline ``$ref`` pointers so ``WithJsonSchema`` never emits orphan refs. - - Workaround for https://github.com/pydantic/pydantic/issues/12145 — - Pydantic's ``GenerateJsonSchema`` does not merge user-provided ``$defs`` - into its internal ref map, so any ``$ref`` in a ``WithJsonSchema`` value - causes a ``KeyError`` when the annotated type is composed into a model. - """ - defs = schema.get("$defs", {}) - - def _resolve(obj): - if isinstance(obj, dict): - if "$ref" in obj: - ref_name = obj["$ref"].split("/")[-1] - if ref_name in defs: - return _resolve(defs[ref_name]) - return {k: _resolve(v) for k, v in obj.items() if k != "$defs"} - if isinstance(obj, list): - return [_resolve(item) for item in obj] - return obj - - return _resolve(schema) - - -@TypeToPydanticType.register(tuple) -def _pydantic_type_tuple(ty): - """Convert finitary tuples to object-based schemas (``properties/required``). - - OpenAI's strict mode rejects the ``prefixItems`` array schema that Pydantic - emits for fixed-length tuples. We convert them to a Pydantic model with - positional ``item_0``, ``item_1``, … fields instead. - - NamedTuples are handled similarly using their field names. - Bare ``tuple`` and variadic ``tuple[T, ...]`` are passed through unchanged. - """ - # NamedTuple subclasses dispatch here via MRO; use field names. - if isinstance(ty, type) and hasattr(ty, "_fields"): - hints = typing.get_type_hints(ty) - nt_fields: list[str] = list(ty._fields) - nt_types = [hints.get(f, typing.Any) for f in nt_fields] - nt_adapters = [pydantic.TypeAdapter(t) for t in nt_types] - nt_model = pydantic.create_model( - ty.__name__, - __config__={"extra": "forbid"}, - **{f: (t, ...) for f, t in zip(nt_fields, nt_types)}, - ) - - def _nt_validate(value, info: pydantic.ValidationInfo): - if isinstance(value, tuple | list): - value = dict(zip(nt_fields, value)) - return ty( - **{ - f: nt_adapters[i].validate_python(value[f], context=info.context) - for i, f in enumerate(nt_fields) - } - ) - - def _nt_serialize(value, info: pydantic.SerializationInfo): - return { - f: nt_adapters[i].dump_python( - getattr(value, f), mode="json", context=info.context - ) - for i, f in enumerate(nt_fields) - } - - return typing.Annotated[ - ty, - pydantic.PlainValidator(_nt_validate), - pydantic.PlainSerializer(_nt_serialize), - pydantic.WithJsonSchema(_inline_refs(nt_model.model_json_schema())), - ] - - args = typing.get_args(ty) - - # Bare tuple or tuple[T, ...] — Pydantic's native handling is fine. - # Note: tuple[()] also has get_args() == (), but has origin=tuple. - if (not args and typing.get_origin(ty) is None) or ( - len(args) == 2 and args[1] is Ellipsis - ): - return ty - - # tuple[()] (empty args with origin) maps to zero fields; otherwise use args. - effective: list[typing.Any] = list(args) - - adapters = [pydantic.TypeAdapter(a) for a in effective] - - model = pydantic.create_model( - "TupleItems", - __config__={"extra": "forbid"}, - **{f"item_{i}": (a, ...) for i, a in enumerate(effective)}, - ) - - def _validate(value, info: pydantic.ValidationInfo): - if isinstance(value, tuple | list): - value = {f"item_{i}": v for i, v in enumerate(value)} - return tuple( - adapters[i].validate_python(value[f"item_{i}"], context=info.context) - for i in range(len(effective)) - ) - - def _serialize(value, info: pydantic.SerializationInfo): - return { - f"item_{i}": adapters[i].dump_python(v, mode="json", context=info.context) - for i, v in enumerate(value) - } - - return typing.Annotated[ - ty, - pydantic.PlainValidator(_validate), - pydantic.PlainSerializer(_serialize), - pydantic.WithJsonSchema(_inline_refs(model.model_json_schema())), - ] - - -@TypeToPydanticType.register(Term) -def _pydantic_type_term(ty: type[Term]): - raise TypeError("Terms cannot be converted to Pydantic types.") - - -@TypeToPydanticType.register(Operation) -def _pydantic_type_operation(ty: type[Operation]): - raise TypeError("Operations cannot be converted to Pydantic types.") - - -@pydantic.validate_call(validate_return=False) -def _validate_image(value: ChatCompletionImageObject) -> Image.Image: - value = pydantic.TypeAdapter(ChatCompletionImageObject).validate_python(value) - image_url: litellm.ChatCompletionImageUrlObject | str = value["image_url"] - url: str = image_url["url"] if isinstance(image_url, dict) else image_url - prefix, data = url.split(",") - if not prefix.startswith("data:image/"): - raise ValueError(f"expected base64 encoded image as data uri, received {url}") - return Image.open(fp=io.BytesIO(base64.b64decode(data))) - - -def _serialize_image(value: Image.Image) -> ChatCompletionImageObject: - buf = io.BytesIO() - value.save(buf, format="PNG") - url = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" - return pydantic.TypeAdapter(ChatCompletionImageObject).validate_python( - {"type": "image_url", "image_url": {"detail": "auto", "url": url}} - ) - - -@TypeToPydanticType.register(Image.Image) -def _pydantic_type_image(ty: type[Image.Image]): - adapter = pydantic.TypeAdapter(ChatCompletionImageObject) - return typing.Annotated[ - ty, - pydantic.PlainValidator(_validate_image), - pydantic.PlainSerializer(_serialize_image), - pydantic.WithJsonSchema(_inline_refs(adapter.json_schema())), - ] - - -class SynthesizedFunction(pydantic.BaseModel): - """Structured output for function synthesis. - - Pydantic model representing synthesized code with function name and module code. - """ - - module_code: str = pydantic.Field( - ..., - description="Complete Python module code (no imports needed)", - ) - - -def _create_typed_synthesized_function( - callable_type: type[Callable], -) -> type[SynthesizedFunction]: - """Create a SynthesizedFunction subclass with type signature in the model description. - - Uses pydantic.create_model to ensure the description is included in the JSON schema - sent to the LLM, informing it of the expected function signature. - """ - if not typing.get_args(callable_type): - type_signature = "Callable" - # Callable[[arg1, arg2, ...], return_type] - elif len(typing.get_args(callable_type)) >= 2: - param_types = typing.get_args(callable_type)[0] - return_type = typing.get_args(callable_type)[-1] - - if param_types is ...: - params_str = "..." - elif isinstance(param_types, list | tuple): - params_str = ", ".join(getattr(t, "__name__", str(t)) for t in param_types) - else: - params_str = str(param_types) - - return_str = getattr(return_type, "__name__", str(return_type)) - type_signature = f"Callable[[{params_str}], {return_str}]" - else: - type_signature = str(callable_type) - - description = f"""Given the specification above, generate a Python function satisfying the following specification and type signature. - -{type_signature} - - -1. Produce one block of Python code. -2. The function MUST have type annotations for all parameters and the return type. -3. The function definition must be the LAST statement - do not add any code after it. -4. Do not include usage examples or function calls. - -""" - - # Use pydantic.create_model to create a proper model with the description - # The __doc__ becomes the model's description in the JSON schema - model = pydantic.create_model( - "TypedSynthesizedFunction", - __base__=SynthesizedFunction, - __doc__=description, - ) - return model - - -def _validate_signature_ast( - func_ast: ast.FunctionDef | ast.AsyncFunctionDef, - expected_params: list[type] | None, -) -> None: - """Validate the function signature from AST before execution.""" - if expected_params is not None: - ast_params = func_ast.args.args + func_ast.args.posonlyargs - if len(ast_params) != len(expected_params): - params_str = ", ".join( - getattr(t, "__name__", str(t)) for t in expected_params - ) - raise ValueError( - f"synthesized function must take exactly {len(expected_params)} " - f"parameter(s) ({params_str}), but got {len(ast_params)}" - ) - - -def _validate_signature_callable( - func: Callable, - expected_params: list[type] | None, - expected_return: type, -) -> None: - """Validate the function signature from runtime callable after execution. - - The synthesized function must have type annotations for parameters and return type. - """ - sig = inspect.signature(func) - - if expected_params is not None: - actual_params = list(sig.parameters.values()) - if len(actual_params) != len(expected_params): - params_str = ", ".join( - getattr(t, "__name__", str(t)) for t in expected_params - ) - return_str = getattr(expected_return, "__name__", str(expected_return)) - raise ValueError( - f"synthesized function must match Callable[[{params_str}], {return_str}] " - f"-- exactly {len(expected_params)} parameter(s) -- " - f"but got {len(actual_params)}" - ) - - actual_return = sig.return_annotation - if actual_return is inspect.Parameter.empty: - raise ValueError( - "decode() requires synthesized function to have a return type annotation" - ) - - -@TypeToPydanticType.register(Callable) -def _pydantic_callable(callable_type: Any) -> Any: - """Create a Pydantic-compatible Annotated type for a parameterized Callable. - - Usage: PydanticCallable(Callable[[int, str], bool]) - """ - type_args = typing.get_args(callable_type) - - if not type_args: - typed_enc = _create_typed_synthesized_function(Callable[..., typing.Any]) # type: ignore[arg-type] - expected_params = None - expected_return = None - else: - if len(type_args) < 2: - raise TypeError( - f"Callable type signature incomplete: {callable_type}. " - "Expected Callable[[ParamTypes...], ReturnType] or Callable[..., ReturnType]." - ) - param_types, expected_return = type_args[0], type_args[-1] - typed_enc = _create_typed_synthesized_function(callable_type) - if param_types is not ... and isinstance(param_types, list | tuple): - expected_params = list(param_types) - else: - expected_params = None - - def _validate(value: Any, info: pydantic.ValidationInfo) -> Callable: - if callable(value) and not isinstance(value, dict): - return value - if isinstance(value, SynthesizedFunction): - encoded = value - elif isinstance(value, dict): - encoded = typed_enc.model_validate(value) - elif isinstance(value, str): - encoded = typed_enc.model_validate_json(value) - else: - raise ValueError( - f"Expected callable, SynthesizedFunction dict, or JSON string, " - f"got {type(value)}" - ) - - if expected_return is None: - raise TypeError( - "Cannot decode/synthesize callable without a concrete type signature. " - "Use Callable[[ParamTypes...], ReturnType] or Callable[..., ReturnType] " - "with a concrete return type (not Any)." - ) - - ctx = info.context or {} - filename = f"" - module: ast.AST = evaluation.parse(encoded.module_code, filename) - - if not isinstance(module, ast.Module) or not module.body: - raise ValueError( - "decode() requires module code with at least one statement." - ) - - last_stmt = module.body[-1] - if not isinstance(last_stmt, ast.FunctionDef): - raise ValueError( - f"decode() requires the last statement to be a function definition, " - f"got {type(last_stmt).__name__}" - ) - - _validate_signature_ast(last_stmt, expected_params) - - # The anchor (Template's underlying function) rides in the decoding context - # under TYPE_CHECK_ANCHOR_KEY; absent for tool-argument decoding, whose - # synthesized Callables are contracted by the tool param's type, not the - # Template's return type, so the Template anchor doesn't apply. When - # present, the code is spliced into the Template body, so first reject - # constructs illegal once nested (star / `__future__` imports), then check. - anchor = ctx.get(TYPE_CHECK_ANCHOR_KEY) - if anchor is not None: - evaluation.scan_non_nestable(module) - spliced = evaluation.splice_into_source(module, anchor) - if spliced is not None: - evaluation.type_check(*spliced) - - g: MutableMapping[str, Any] = {} - g.update( - { - k: v - for k, v in ctx.items() - if k.isidentifier() and k != TYPE_CHECK_ANCHOR_KEY - } - ) - bytecode: types.CodeType = evaluation.compile(module, filename) - evaluation.exec(bytecode, g) - - func_name = last_stmt.name - if func_name not in g: - raise ValueError( - f"decode() expected function '{func_name}' to be defined in globals" - ) - - result = g[func_name] - if not callable(result): - raise ValueError( - f"decode() expected '{func_name}' to be callable, got {type(result)}" - ) - - _validate_signature_callable(result, expected_params, expected_return) - return result - - def _serialize(value: Callable) -> dict: - if not callable(value): - raise TypeError(f"Expected callable, got {type(value)}") - - try: - source = inspect.getsource(value) - except (OSError, TypeError): - source = None - - if source: - return typed_enc(module_code=textwrap.dedent(source)).model_dump() - - name = getattr(value, "__name__", None) - docstring = inspect.getdoc(value) - if name is None or docstring is None: - raise ValueError( - f"Cannot encode callable {value}: no source code and no __name__ or docstring" - ) - - try: - sig = inspect.signature(value) - sig_str = str(sig) - except (ValueError, TypeError): - sig_str = "(...)" - - stub_code = f'''def {name}{sig_str}: - """{docstring}""" - ... -''' - return typed_enc(module_code=stub_code).model_dump() - - return typing.Annotated[ - callable_type, - pydantic.PlainValidator(_validate), - pydantic.PlainSerializer(_serialize), - pydantic.WithJsonSchema( - _inline_refs(pydantic.TypeAdapter(typed_enc).json_schema()) - ), - ] - - -def _validate_tool( - value: ChatCompletionToolParam, info: pydantic.ValidationInfo -) -> Tool: - assert isinstance(info.context, Mapping), "Tool decoding requires context" - value = pydantic.TypeAdapter(ChatCompletionToolParam).validate_python(value) - try: - return info.context[value["function"]["name"]] - except KeyError as e: - raise NotImplementedError(f"Unknown tool: {value['function']['name']}") from e - - -def _serialize_tool( - value: Tool, info: pydantic.SerializationInfo -) -> ChatCompletionToolParam: - fields: dict[str, Any] = { - name: TypeToPydanticType().evaluate(param.annotation) - for name, param in inspect.signature(value).parameters.items() - } - sig_model = pydantic.create_model( - "Params", - __config__={"extra": "forbid"}, - **fields, - ) - response_format = litellm.utils.type_to_response_format_param(sig_model) - assert response_format is not None - assert value.__default__.__doc__ is not None - # Advertise under the context key, since decode (`_validate_tool`) resolves the call by that name. - tool_name = value.__name__ - context = info.context - if isinstance(context, Mapping): - for key, tool in context.items(): - if tool is value: - tool_name = key - break - return pydantic.TypeAdapter(ChatCompletionToolParam).validate_python( - { - "type": "function", - "function": { - "name": tool_name, - "description": textwrap.dedent(value.__default__.__doc__), - "parameters": response_format["json_schema"]["schema"], - "strict": True, - }, - } - ) - - -@TypeToPydanticType.register(Tool) -def _pydantic_type_tool(ty: type[Tool]): - schema = _inline_refs(pydantic.TypeAdapter(ChatCompletionToolParam).json_schema()) - schema = _ensure_strict_json_schema(schema, path=(), root={}) - return typing.Annotated[ - ty, - pydantic.PlainValidator(_validate_tool), - pydantic.PlainSerializer(_serialize_tool), - pydantic.WithJsonSchema(schema), - ] - - -def _validate_tool_call( - value: ChatCompletionMessageToolCall, - info: pydantic.ValidationInfo, -) -> DecodedToolCall: - if isinstance(value, dict): - value = OpenAIChatCompletionMessageToolCall.model_validate(value) - ctx = info.context or {} - assert value.function.name is not None - tool = ctx[value.function.name] - assert isinstance(tool, Tool) - sig = inspect.signature(tool) - decoded_args = {} - for name, raw_arg in json.loads(value.function.arguments).items(): - assert name in sig.parameters, ( - f"Unexpected argument {name} for tool {tool.__name__}" - ) - param = sig.parameters[name] - arg_enc: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter( - Encodable[param.annotation] # type: ignore[name-defined] - ) - decoded_args[name] = arg_enc.validate_python(raw_arg, context=ctx) - return DecodedToolCall( - tool=tool, - bound_args=sig.bind(**decoded_args), - id=value.id, - name=value.function.name, - ) - - -def _serialize_tool_call( - value: DecodedToolCall, info: pydantic.SerializationInfo -) -> dict: - ctx = info.context or {} - encoded_args = {} - for k, v in value.bound_args.arguments.items(): - v_enc: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter( - Encodable[nested_type(v).value] # type: ignore[misc] - ) - encoded_args[k] = v_enc.dump_python(v, mode="json", context=ctx) - return OpenAIChatCompletionMessageToolCall.model_validate( - { - "type": "function", - "id": value.id, - "function": { - "name": value.tool.__name__, - "arguments": json.dumps(encoded_args), - }, - } - ).model_dump(mode="json") - - -@TypeToPydanticType.register(DecodedToolCall) -def _pydantic_type_tool_call(ty: type[DecodedToolCall]): - # Use OpenAI's ChatCompletionMessageToolCall (has actual fields: id, function, - # type) rather than litellm's (empty dict with extra="allow"). - schema = _inline_refs(OpenAIChatCompletionMessageToolCall.model_json_schema()) - schema = _ensure_strict_json_schema(schema, path=(), root={}) - return typing.Annotated[ - ty, - pydantic.PlainValidator(_validate_tool_call), - pydantic.PlainSerializer(_serialize_tool_call), - pydantic.WithJsonSchema(schema), - ] diff --git a/effectful/handlers/llm/evaluation.py b/effectful/handlers/llm/evaluation.py deleted file mode 100644 index 729d8185c..000000000 --- a/effectful/handlers/llm/evaluation.py +++ /dev/null @@ -1,705 +0,0 @@ -import ast -import builtins -import code -import codeop -import collections.abc -import contextlib -import inspect -import io -import json -import linecache -import logging -import os -import shutil -import subprocess -import sys -import tempfile -import typing -from collections.abc import MutableMapping -from types import CodeType -from typing import Any - -from RestrictedPython import ( - Eval, - Guards, - RestrictingNodeTransformer, - compile_restricted, - safe_globals, -) -from RestrictedPython.PrintCollector import PrintCollector - -from effectful.handlers.llm.template import Tool -from effectful.ops.syntax import ObjectInterpretation, defop, implements -from effectful.ops.types import Operation - - -@defop -def parse(source: str, filename: str) -> ast.Module: - """ - Parse source text into an AST. - - source: The Python source code to parse. - filename: The filename recorded in the resulting AST for tracebacks and tooling. - - Returns the parsed AST. - """ - raise NotImplementedError( - "An eval provider must be installed in order to parse code." - ) - - -@defop -def type_check( - source: str, - lo: int | None = None, - hi: int | None = None, - *, - lenient: bool = False, -) -> None: - """ - Type check a module source, reporting only diagnostics inside a line region. - - source: A complete module source to check (e.g. produced by - ``splice_into_source``, which splices generated code into a Template's real - module source). - lo, hi: Inclusive line range within ``source`` to report errors from; when - omitted, the whole source is in scope. Errors outside the region are - ignored so unrelated pre-existing code never blocks synthesis. - lenient: when True, relax mypy for incrementally-built REPL code spliced into a - Template body -- allow redefinition (a cell may rebind or redefine a name) - and don't require the body to satisfy the Template's return type. Off (strict) - for synthesized ``Callable`` bodies, which must honor their signature. - - Returns None, raises TypeError on an in-region failure. - """ - raise NotImplementedError( - "An eval provider must be installed in order to type check code." - ) - - -@defop -def compile(module: ast.Module, filename: str) -> CodeType: - """ - Compile an AST into a Python code object. - - module: The AST to compile (typically produced by parse()). - filename: The filename recorded in the resulting code object (CodeType.co_filename), used in tracebacks and by inspect.getsource(). - - Returns the compiled code object. - """ - raise NotImplementedError( - "An eval provider must be installed in order to compile code." - ) - - -@defop -def exec( - bytecode: CodeType, - env: dict[str, Any], -) -> None: - """ - Execute a compiled code object. - - bytecode: A code object to execute (typically produced by compile()). - env: The namespace mapping used during execution. - - After ``exec(bytecode, env)`` returns, ``env`` reflects all top-level - binding effects of the executed code (new names and rebindings alike). - """ - raise NotImplementedError( - "An eval provider must be installed in order to execute code." - ) - - -logger = logging.getLogger(__name__) - - -def scan_non_nestable(generated: ast.Module) -> None: - """Reject constructs legal at module level but illegal once nested in a function. - - ``from ... import *`` and ``from __future__ import ...`` are both ``SyntaxError``s - inside a function body, but mypy *accepts* a nested star import silently, so the - splice would slip an illegal construct past the type check and fail later at - ``compile``/``exec``. Detect them explicitly and raise before splicing. Raises - ``ValueError`` (this is rejecting invalid generated *source*, not signaling a type - error), so a decoder can catch it alongside ``SyntaxError`` without swallowing a real - ``TypeError`` from a broken provider. - """ - for stmt in generated.body: - if isinstance(stmt, ast.ImportFrom): - if stmt.module == "__future__": - raise ValueError( - "generated code uses `from __future__ import ...`, which is " - "illegal once spliced into a function body" - ) - if any(alias.name == "*" for alias in stmt.names): - raise ValueError( - "generated code uses a star import (`from ... import *`), which " - "is illegal once spliced into a function body" - ) - - -def _def_nodes( - module: ast.Module, -) -> list[ast.FunctionDef | ast.AsyncFunctionDef]: - """All function definitions in ``module``, in a stable order that an - ``ast.unparse`` -> ``ast.parse`` round-trip preserves (so a def keeps its - index across it).""" - return [ - n - for n in ast.walk(module) - if isinstance(n, ast.FunctionDef | ast.AsyncFunctionDef) - ] - - -def _find_def_at_lineno( - module: ast.Module, lineno: int -) -> ast.FunctionDef | ast.AsyncFunctionDef | None: - """Locate the function definition whose definition site is ``lineno``. - - Matches ``fn.__code__.co_firstlineno`` -- the first decorator line, or the - ``def`` line when undecorated -- which identifies the def directly and - unambiguously (no name matching, and nesting-agnostic). Returns None only if - no def starts there: a dynamically generated ``fn`` with no source def, or - source that has drifted since import. - """ - for node in _def_nodes(module): - start = node.decorator_list[0].lineno if node.decorator_list else node.lineno - if start == lineno: - return node - return None - - -def _region_errors(stdout: str, lo: int | None, hi: int | None) -> list[dict[str, Any]]: - """mypy ``--output=json`` diagnostics of severity ``error`` whose reported - line falls within ``[lo, hi]`` -- the spliced region. An open bound (``None``) - is unbounded on that side, so ``lo=hi=None`` reports every error. - - ``--output=json`` emits one JSON object per diagnostic carrying mypy's own - ``severity`` and ``line`` fields, so we filter on those directly rather than - parsing (and risking mis-parsing) its human-readable format. Only reached - for exit status < 2; a fatal status emits text, not JSON, and is handled by - the caller before this runs. - """ - errors: list[dict[str, Any]] = [] - for line in stdout.splitlines(): - if not line.strip(): - continue - diag = json.loads(line) - if diag["severity"] != "error": - continue - if (lo is None or lo <= diag["line"]) and (hi is None or diag["line"] <= hi): - errors.append(diag) - return errors - - -def splice_into_source( - generated: ast.Module, anchor: Any -) -> tuple[str, int, int] | None: - """Splice `generated` into the anchor Template's own function body, in its real - module source. - - Returns the modified module source and the ``[lo, hi]`` line span of the - spliced body within it, or ``None`` when the anchor's source can't be recovered - (the caller skips rather than guesses). Raises ``RuntimeError`` if the source is - recovered but the anchor's def can't be located in it (source drift) -- a real - error, not a silent pass. - - The generated function -- and any helpers it defines alongside -- becomes the - body of the Template's own function at its real (possibly nested) position, so - the generated code is checked in its real lexical scope with no synthesized - type stubs. - """ - if not generated.body: - raise TypeError("splice: generated module is empty") - last = generated.body[-1] - if not isinstance(last, ast.FunctionDef | ast.AsyncFunctionDef): - raise TypeError( - f"splice: last statement must be a function definition, " - f"got {type(last).__name__}" - ) - target_name = last.name - - recovered = _recover_template_def(anchor) - if recovered is None: - return None - module_ast, template_def = recovered - - # Splice in place: replace the body with the generated body and bind the - # target against the (source) return annotation via `return`. Decorators are - # left untouched -- mypy checks a function's body against its declared return - # type regardless of decorators (even an unresolvable / `Any` one), and the - # decorator application itself doesn't spuriously fail, so touching the - # surrounding source as little as possible keeps the splice robust. - template_def.body = [ - *generated.body, - ast.Return(ast.Name(target_name, ast.Load())), - ] - - # mypy reports line numbers in the coordinates of `checked_source`, so we need - # the spliced *body's* span there. ast.unparse reassigns line numbers but - # preserves def order, so the def keeps its index in walk order -- take the def - # at that same index in the re-parsed source. - # - # The region is the body (the generated code) only, NOT the def header: the - # signature and decorators are the Template author's own pre-existing source, - # which we must not attribute to synthesis. This matters for templates whose - # module source can't be fully recovered -- notably notebook/REPL cells, which - # share a runtime namespace but whose recovered source is a single cell missing - # the other cells' imports, so the signature's own annotations (e.g. `Literal`, - # `Callable`) look undefined to mypy. Flagging only the body keeps those - # spurious signature-line diagnostics out of the gate. - def_index = _def_nodes(module_ast).index(template_def) - checked_source = ast.unparse(ast.fix_missing_locations(module_ast)) - spliced = _def_nodes(ast.parse(checked_source))[def_index] - lo = spliced.body[0].lineno # first generated statement (body is non-empty) - hi = spliced.end_lineno or lo - return checked_source, lo, hi - - -def _recover_template_def( - anchor: Any, -) -> tuple[ast.Module, ast.FunctionDef | ast.AsyncFunctionDef] | None: - """Locate the anchor Template's own ``def`` in its real module source. - - Returns the parsed module AST and the def node, or ``None`` when the source can't - be recovered (REPL/exec/notebook Template with no linecache entry -- the caller - skips rather than guesses). Raises ``RuntimeError`` on source drift (source - recovered but the def no longer sits where ``fn`` was compiled from). - """ - fn = inspect.unwrap(anchor) # staticmethod/classmethod -> underlying function - # Recover the module source via fn's own filename -- a real path or a - # linecache-registered synthetic name (e.g. ) for REPL/exec/ - # notebook templates; linecache.getlines reads real files from disk too. - try: - source_file = inspect.getsourcefile(fn) - except TypeError: - source_file = None - module_source = "".join(linecache.getlines(source_file)) if source_file else "" - if not module_source: - logger.warning("skipping type check: cannot recover source for %r", fn) - return None - module_ast = ast.parse(module_source) - template_def = _find_def_at_lineno(module_ast, fn.__code__.co_firstlineno) - if template_def is None: - raise RuntimeError( - f"cannot locate {getattr(fn, '__qualname__', fn)!r} in its module " - f"source (source drifted since import?)" - ) - return module_ast, template_def - - -def _splice_repl( - prior: list[str], snippet: str, anchor: Any -) -> tuple[str, int, int] | None: - """Splice the cumulative REPL code -- ``prior`` snippets followed by the current - ``snippet`` -- into the anchor Template's body, in its real module source, and return - the modified source with the ``[lo, hi]`` line span of the *current* snippet. - - The REPL code becomes the Template function's body at its real (possibly nested) - position, so the Template's parameters and enclosing scope -- i.e. the session's seed - env -- are in scope, and each snippet sees the ones before it (they are function - locals). No ``return`` is appended; the REPL code doesn't produce the Template's - declared type, and that contract is waived by ``lenient`` type checking. Every prior - snippet stays in the body so its bindings resolve (matching the runtime, which ran - them), but only the current snippet's lines are reported, so an earlier cell's error - isn't re-reported on every later call. - - Returns ``None`` when the current snippet has no statements to check, or when the - Template's source can't be recovered -- a Template defined at a REPL, in a notebook, or - via ``exec()`` is sourceless, so we skip the check and run the code unchecked, exactly - as ``splice_into_source`` does for a sourceless Callable anchor. Raises ``RuntimeError`` - only on source *drift* (source recovered but the def no longer sits where it was - compiled from), which ``_recover_template_def`` surfaces. - """ - # An empty or comment-only snippet parses to zero statements: nothing to check. - n_current = len(ast.parse(snippet).body) - if n_current == 0: - return None - # None means the Template's source can't be recovered (REPL/exec/notebook-defined) -- - # skip, like the Callable path, rather than break the tool; `_recover_template_def` - # raises on source drift, which is a real error and propagates. - recovered = _recover_template_def(anchor) - if recovered is None: - return None - module_ast, template_def = recovered - cumulative = "".join(s if s.endswith("\n") else s + "\n" for s in [*prior, snippet]) - template_def.body = ast.parse(cumulative).body - - # mypy reports line numbers in the coordinates of the unparsed source; the current - # snippet is the last `n_current` statements of the spliced body. ast.unparse keeps def - # order, so the template def is at the same walk index after the round-trip. - def_index = _def_nodes(module_ast).index(template_def) - checked_source = ast.unparse(ast.fix_missing_locations(module_ast)) - spliced = _def_nodes(ast.parse(checked_source))[def_index] - lo = spliced.body[-n_current].lineno - hi = spliced.body[-1].end_lineno or lo - return checked_source, lo, hi - - -def _mypy_check_region( - source: str, - lo: int | None = None, - hi: int | None = None, - lenient: bool = False, -) -> None: - """Run mypy on `source` and raise ``TypeError`` if any error diagnostic falls - within ``[lo, hi]``; raise ``RuntimeError`` if mypy itself fails to run. - - Applies mypy to whatever source it's given -- spliced or otherwise -- and - reports only the region's errors (the whole source when the region is - omitted), so pre-existing errors elsewhere in `source` never block synthesis. - - When ``lenient`` (for REPL code spliced into a Template body): allow a variable to be - redefined with a new type across cells (``--allow-redefinition``), a def/class/import - to be redefined (``no-redef``), and the body not to return the Template's declared type - (``return``/``empty-body``). All normal for an incrementally-built REPL, not real errors. - """ - lenient_flags = ( - [ - "--allow-redefinition", - "--disable-error-code=no-redef", - "--disable-error-code=return", - "--disable-error-code=empty-body", - ] - if lenient - else [] - ) - # Run mypy as a subprocess, not the in-process `mypy.api.run`: the API builds - # typeshed and a full module graph inside this process and never returns that - # memory, so under a test/agent session doing many checks it accumulates to many - # GB (OOM). A subprocess reclaims all of it on exit. Pass a file (not --command: - # it hits an argv length limit on large modules); each call gets an isolated temp - # dir + cache so parallel decodes don't share -- and deadlock on -- mypy's cache. - tmpdir = tempfile.mkdtemp(prefix="effectful_typecheck_") - try: - tf_path = os.path.join(tmpdir, "_synthesized.py") - with open(tf_path, "w", encoding="utf-8") as f: - f.write(source) - proc = subprocess.run( - [ - sys.executable, - "-m", - "mypy", - tf_path, - "--cache-dir", - os.path.join(tmpdir, "cache"), - "--no-error-summary", - "--output=json", - "--ignore-missing-imports", - "--disable-error-code=import-untyped", - *lenient_flags, - ], - capture_output=True, - text=True, - ) - stdout, stderr, status = proc.stdout, proc.stderr, proc.returncode - finally: - shutil.rmtree(tmpdir, ignore_errors=True) - # Exit status >= 2 means mypy itself failed (fatal/usage/internal/syntax) -- a - # tool failure, not a type error -- and it emits text rather than JSON, so - # raise `RuntimeError` rather than parse or silently pass. - if status >= 2: - raise RuntimeError( - f"mypy could not check the source:\n{(stdout or '') + (stderr or '')}" - ) - errors = _region_errors(stdout or "", lo, hi) - if errors: - # Not the source: it's large and the model already has the generated code. - report = "\n".join(json.dumps(e) for e in errors) - raise TypeError("mypy type check failed:\n" + report) - - -# Eval Providers - - -class UnsafeEvalProvider(ObjectInterpretation): - """UNSAFE provider that handles parse, comple and exec operations - by shelling out to python *without* any further checks. Only use for testing.""" - - @implements(type_check) - def type_check( - self, - source: str, - lo: int | None = None, - hi: int | None = None, - *, - lenient: bool = False, - ) -> None: - _mypy_check_region(source, lo, hi, lenient) - - @implements(parse) - def parse(self, source: str, filename: str) -> ast.Module: - # Cache source under `filename` so inspect.getsource() can retrieve it later. - # inspect uses f.__code__.co_filename -> linecache.getlines(filename) - linecache.cache[filename] = ( - len(source), - None, - source.splitlines(True), - filename, - ) - - return ast.parse(source, filename=filename, mode="exec") - - @implements(compile) - def compile(self, module: ast.AST, filename: str) -> CodeType: - return builtins.compile(typing.cast(typing.Any, module), filename, "exec") - - @implements(exec) - def exec( - self, - bytecode: CodeType, - env: dict[str, Any], - ) -> None: - # Ensure builtins exist in the execution environment. - env.setdefault("__builtins__", __builtins__) - - # Execute module-style so top-level defs land in `env`. - builtins.exec(bytecode, env, env) - - -class _StdoutPrintCollector(PrintCollector): - """`_print_` factory whose `print(...)` writes to the real `sys.stdout` - (so output-capturing callers see it) rather than accumulating into the - collector's discarded `printed` buffer.""" - - def _call_print(self, *objects, **kwargs): - kwargs.setdefault("file", sys.stdout) - builtins.print(*objects, **kwargs) - - -class RestrictedEvalProvider(ObjectInterpretation): - """ - Safer provider using RestrictedPython. - - RestrictedPython is not a complete sandbox, but it enforces a restricted - language subset and expects you to provide a constrained exec environment. - - policy : dict[str, Any], optional - RestrictedPython compile_restricted policy for compilation - """ - - policy: type[RestrictingNodeTransformer] | None = None - - def __init__( - self, - *, - policy: type[RestrictingNodeTransformer] | None = None, - ): - self.policy = policy - - @implements(type_check) - def type_check( - self, - source: str, - lo: int | None = None, - hi: int | None = None, - *, - lenient: bool = False, - ) -> None: - _mypy_check_region(source, lo, hi, lenient) - - @implements(parse) - def parse(self, source: str, filename: str) -> ast.Module: - # Keep inspect.getsource() working for dynamically-defined objects. - linecache.cache[filename] = ( - len(source), - None, - source.splitlines(True), - filename, - ) - return ast.parse(source, filename=filename, mode="exec") - - @implements(compile) - def compile(self, module: ast.Module, filename: str) -> CodeType: - # RestrictedPython can compile from an AST directly. - return compile_restricted( - module, - filename=filename, - mode="exec", - policy=self.policy or RestrictingNodeTransformer, - ) - - @implements(exec) - def exec( - self, - bytecode: CodeType, - env: dict[str, Any], - ) -> None: - # Build restricted globals from RestrictedPython's defaults - rglobals: dict[str, Any] = safe_globals.copy() - - # Enable class definitions (required for Python 3) - rglobals["__metaclass__"] = type - rglobals["__name__"] = "restricted" - - # Layer `env` on top (without letting callers replace the restricted builtins). - rglobals.update({k: v for k, v in env.items() if k != "__builtins__"}) - - # Enable for loops and comprehensions - rglobals["_getiter_"] = Eval.default_guarded_getiter - # Enable sequence unpacking in comprehensions and for loops - rglobals["_iter_unpack_sequence_"] = Guards.guarded_iter_unpack_sequence - - rglobals["getattr"] = Guards.safer_getattr - rglobals["setattr"] = Guards.guarded_setattr - rglobals["_write_"] = lambda x: x - - # RestrictedPython rewrites `print(...)` into its `_print_` collector - # protocol; route it to the real stdout so output-capturing callers - # (e.g. redirect_stdout) see it instead of a discarded collector. - rglobals["_print_"] = _StdoutPrintCollector - - # Snapshot value identities before execution so we can copy back every - # *binding effect* — both new names and rebindings of seeded names. - before = dict(rglobals) - builtins.exec(bytecode, rglobals, rglobals) - - sentinel = object() - env.update( - { - key: value - for key, value in rglobals.items() - if key != "__builtins__" and before.get(key, sentinel) is not value - } - ) - - -class _OpCommandCompiler(codeop.CommandCompiler): - """A `codeop.CommandCompiler` that routes compilation through the - `parse`/`compile` effect operations (so the installed eval provider owns it - and `parse` populates `linecache`), replacing the native single-mode - compiler that `code.InteractiveInterpreter` installs. - """ - - def __call__( - self, source: str, filename: str = "", symbol: str = "single" - ) -> CodeType: - # `runsource` passes symbol="single"; we ignore it and compile in the - # exec mode the ops produce, so a complete multi-statement block runs in - # one shot. Incomplete/invalid input raises SyntaxError, which - # `runsource` routes to `showsyntaxerror` (we do not buffer partial input - # -- there is no line-at-a-time protocol). - return compile(parse(source, filename), filename) - - -class ReplSession(code.InteractiveInterpreter): - """A persistent, output-capturing Python session seeded from a lexical - context. - - `exec_code(source)` runs a pre-compiled code object in `self.locals` through - the `exec` effect operation. Both bindings and captured stdout/stderr - persist across calls -- variables, imports and definitions accumulate exactly - like a REPL -- and the session (with its buffer) is discarded as a whole when - it goes out of scope. Each call returns only the output it produced; a - snippet that raises has its traceback appended to that output rather than - propagating -- mirroring `code.InteractiveInterpreter`, only `SystemExit` - propagates -- so failures are surfaced as text. There is no bare-expression - auto-echo, so use `print()` to surface values. - - Compilation -- and therefore syntax checking -- happens earlier, at the - `Encodable[CodeType]` boundary; this session only executes. - """ - - # The session's captured output, accumulated across calls and exposed for - # introspection. stdout (`print` output) and stderr (writes plus tracebacks) - # are kept separate; `exec_code` returns each call's slice of both. - stdout: io.StringIO - stderr: io.StringIO - - def __init__(self, env: MutableMapping[str, Any]): - # Run in a fresh writable dict seeded with a flat view of `env`. This is - # forced by `exec`: its globals must be one real dict (a ChainMap is - # rejected), and a REPL needs a single persistent namespace so a function - # defined in one snippet sees a name a later snippet binds. Seeding a flat - # copy also leaves the lexical seed untouched, so REPL assignments never - # leak into the surrounding scope. - scope: dict[str, Any] = dict(env) - # When `env` is the per-call `ChainMap` (its outer layers are read-only - # frame proxies), splice this dict in as an extra shadowing first layer so - # the bindings are *also* visible to the rest of the Template call - # (mirroring `exec`) -- still scoped to the call, since that ChainMap is. - if isinstance(env, collections.ChainMap): - env.maps.insert(0, scope) - # `InteractiveInterpreter.__init__` stores it as `self.locals`, so we reuse - # the base's runcode/showtraceback/write machinery. - super().__init__(scope) - # Route `runsource`'s compilation through the `parse`/`compile` ops too, so - # it stays consistent with our `runcode` (which execs through the `exec` - # op) rather than the native single-mode compiler the base installed. - self.compile = _OpCommandCompiler() - self.stdout = io.StringIO() - self.stderr = io.StringIO() - self._prior_snippets: list[str] = [] - - @property - def prior_snippets(self) -> list[str]: - """Sources of the actual error-free executed snippets, in order -- the type-check - context the `Encodable[CodeType]` decoder splices before the current snippet.""" - return self._prior_snippets - - def runcode(self, code: CodeType) -> None: - # Mirrors `InteractiveInterpreter.runcode` exactly; the only difference - # is that `exec` here is the effect operation, so execution routes - # through the installed eval provider. `showtraceback` reports failures - # via `self.write`, which `exec_code` has redirected into `self.stderr`. - try: - exec(code, self.locals) - except SystemExit: - raise - except: - self.showtraceback() - - @Tool.define - def exec_code(self, code: CodeType) -> str: - """Run Python in a persistent, stateful session and return its output. - - This is a long-lived REPL, not a one-shot sandbox: every call runs in the - SAME namespace, so names you bind in one call stay available in later - calls within the same task. Imports, function/class definitions and - variable assignments all accumulate during the session of this template. - The namespace starts seeded with the in-scope variables of the surrounding context, which you may read and - rebind. - - Output: returns this call's output -- its stdout (what `print` wrote) - followed by its stderr (which includes the traceback if the code raised). - There is NO automatic echoing of results -- a bare expression on its own - line (e.g. `1 + 1`) displays nothing, so call `print(...)` for anything - you want to see. A snippet that raises has its traceback returned and the - session survives, so you can read the error and continue in the next call - (only `SystemExit` aborts). - - Provide `code` as a string of Python source. It must be a complete, - compilable snippet -- incomplete or invalid source is rejected before it - runs. - """ - out_start = self.stdout.tell() - err_start = self.stderr.tell() - # Record this snippet's source so the *next* snippet's decode-time type check can - # splice the accumulated session code into the Template body. The type check itself - # lives in the `Encodable[CodeType]` decoder (as it does for synthesized Callables), - # not here -- this session only runs code. - self._prior_snippets.append("".join(linecache.getlines(code.co_filename))) - with ( - contextlib.redirect_stdout(self.stdout), - contextlib.redirect_stderr(self.stderr), - ): - self.runcode(code) - return self.stdout.getvalue()[out_start:] + self.stderr.getvalue()[err_start:] - - -@Operation.define -def _repl_session(env: MutableMapping[str, Any]) -> "ReplSession": - """Return the REPL session for the current Template call, seeded from `env`. - - `PythonRepl` (in completions.py) installs a fresh handler for this inside each - `Template.__apply__` (mirroring how `__history__` is managed), giving the session a - lifetime of exactly one Template call. Outside such a scope there is no managed - session, so this falls back to a fresh one -- e.g. when tools are listed outside a - Template call, or when a code object is decoded with no REPL in scope. - - Defined here (not with `PythonRepl`) so the `Encodable[CodeType]` decoder can reach the - session -- and its accumulated `prior_snippets` -- at decode time without importing - `completions` (which would be a cycle). - """ - return ReplSession(env) diff --git a/effectful/handlers/llm/harness/__init__.py b/effectful/handlers/llm/harness/__init__.py new file mode 100644 index 000000000..887119f38 --- /dev/null +++ b/effectful/handlers/llm/harness/__init__.py @@ -0,0 +1,317 @@ +"""A reusable harness for running `effectful.handlers.llm` example scripts. + +The example scripts under ``docs/source/llm_examples`` share a fixed stack of +handlers -- a LiteLLM provider, a Python REPL, retry/decoding logic, and so on -- +that turns a bare `Template`/`Agent` into something runnable. This module +factors that stack into a single object, `harness`, so the scripts themselves +carry none of the boilerplate. + +`harness` is a `contextlib.ContextDecorator`, so it can be used programmatically +either as a context manager or as a decorator:: + + with harness(model="gpt-4o", render=True): + main() + + @harness(model="gpt-4o") + def main() -> None: + ... + +Run as a module it becomes a command-line launcher that wraps an arbitrary +script in the same context:: + + python -m effectful.handlers.llm.harness + +Harness flags (``--model``, ``--num-retries``, ``--langfuse``, ``--render``, +``--dump-system-prompt``, ``--tool-choice``, ``--reasoning-effort``, +``--eval-provider``, ``--pdb``, ``--persist-db``) are consumed here; every other +flag is passed through to the script unchanged. +""" + +import argparse +import contextlib +import inspect +import os +import pathlib +import pdb +import runpy +import sys +import textwrap +import typing + +import litellm +import tenacity + +from effectful.handlers.llm.harness.display import ( + RichTerminalRenderer, + SystemPromptDumper, +) +from effectful.handlers.llm.harness.durability import TenacityRetryer +from effectful.handlers.llm.harness.execution.builtin import BuiltinExecutor +from effectful.handlers.llm.harness.execution.restricted import ( + RestrictedPythonExecutor, +) +from effectful.handlers.llm.harness.observability import LangfuseTracer +from effectful.handlers.llm.harness.persistence import SQLitePersister +from effectful.handlers.llm.harness.provision import ( + LiteLLMProvider, +) +from effectful.handlers.llm.harness.synthesis import ( + FinalBodySynthesizer, + StatefulReplSynthesizer, +) +from effectful.ops.semantics import handler + +# The providers that run model-authored Python, by the name the CLI knows them +# under: `unsafe` runs it with the plain interpreter, `restricted` under +# RestrictedPython's language subset and a guarded environment. +EVAL_PROVIDERS: dict[str, typing.Callable[[], typing.Any]] = { + "unsafe": BuiltinExecutor, + "restricted": RestrictedPythonExecutor, +} + + +class harness(contextlib.ContextDecorator): + """Install the standard `effectful.handlers.llm` handler stack. + + Constructing a `harness` records the configuration; entering it (as a + context manager, decorator, or via the module CLI) installs the handlers and + exiting removes them. The handlers, in installation order, are: + + 1. `LiteLLMProvider` -- the model backend. + 2. `TerminalRenderer` -- live-render the streaming history (if ``render``). + 3. `SystemPromptDumper` -- dump the system prompt (if ``dump_system_prompt``). + 4. The ``eval_provider`` (`EVAL_PROVIDERS`) and `PythonRepl` -- run + model-authored Python. + 5. `SynthesizeAndCall` -- synthesize a function and call it. + 6. `RetryLLMHandler` -- retry malformed/failing model output. + 7. `LexicalReaders` -- expose lexically-scoped tools to the model. + 8. `SQLitePersister` -- checkpoint a persisted `Agent`'s state/history to + SQLite after each successful call (if ``persist_db``). + 9. `LangfuseTracer` -- log calls to Langfuse (if ``langfuse``). + + Args: + model: LLM model to use. + num_retries: Attempts for malformed/failing model output. + langfuse: Log LLM calls and metadata to Langfuse. + render: Live-render the streaming message history in the terminal. + dump_system_prompt: If set, dump the assembled system prompt to this + Markdown file. + tool_choice: ``tool_choice`` forwarded to the provider. + reasoning_effort: ``reasoning_effort`` forwarded to the provider (and on + to ``litellm.completion``); omitted from requests when ``None``. + api_base: API base URL forwarded to the provider. + api_key: API key forwarded to the provider. + persist_db: If set, path to a SQLite database used to checkpoint a + persisted `~effectful.handlers.llm.template.Agent`'s (one + constructed with an explicit `agent_id`) state and history via + `~effectful.handlers.llm.completions.SQLitePersister`. + eval_provider: Which provider runs model-authored Python -- a key of + `EVAL_PROVIDERS` (``"unsafe"`` or ``"restricted"``). + """ + + def __init__( + self, + *, + model: str = "", + num_retries: int = 5, + langfuse: bool = False, + render: bool = False, + dump_system_prompt: str | os.PathLike[str] | None = None, + tool_choice: str = "auto", + reasoning_effort: str | None = None, + api_base: str | None = None, + api_key: str | None = None, + persist_db: str | os.PathLike[str] | None = None, + eval_provider: str = "unsafe", + ) -> None: + self.model = model + self.num_retries = num_retries + self.langfuse = langfuse + self.render = render + self.dump_system_prompt = dump_system_prompt + self.tool_choice = tool_choice + self.reasoning_effort = reasoning_effort + self.api_base = api_base + self.api_key = api_key + self.persist_db = persist_db + self.eval_provider = eval_provider + + def __enter__(self) -> "harness": + stack = contextlib.ExitStack() + # Only forward `reasoning_effort` when set, so we don't send + # `reasoning_effort=None` on every request to providers that reject it. + provider_config: dict[str, str] = {} + if self.reasoning_effort is not None: + provider_config["reasoning_effort"] = self.reasoning_effort + stack.enter_context( + handler( + LiteLLMProvider( + model=self.model, + tool_choice=self.tool_choice, + api_base=self.api_base, + api_key=self.api_key, + **provider_config, + ) + ) + ) + if self.render: + stack.enter_context(handler(RichTerminalRenderer())) + if self.dump_system_prompt: + stack.enter_context( + handler(SystemPromptDumper(path=pathlib.Path(self.dump_system_prompt))) + ) + stack.enter_context(handler(EVAL_PROVIDERS[self.eval_provider]())) + stack.enter_context(handler(StatefulReplSynthesizer())) + stack.enter_context(handler(FinalBodySynthesizer())) + stack.enter_context( + handler(TenacityRetryer(stop=tenacity.stop_after_attempt(self.num_retries))) + ) + # stack.enter_context(handler(LexicalReaders())) + if self.persist_db is not None: + stack.enter_context(handler(SQLitePersister(pathlib.Path(self.persist_db)))) + if self.langfuse: + stack.enter_context(handler(LangfuseTracer())) + self._stack = stack + return self + + def __exit__(self, *exc_info) -> bool | None: + return self._stack.__exit__(*exc_info) + + +def _reasoning_effort_choices() -> list[str] | None: + """The ``reasoning_effort`` values ``litellm.completion`` declares. + + Extracted from the ``Optional[Literal[...]]`` annotation on the live + signature so the CLI choices track litellm exactly across upgrades. Returns + ``None`` (leave the flag unrestricted) if the annotation isn't a Literal we + can read, so a shape change in litellm degrades to accepting any string + rather than breaking the launcher. + """ + try: + annotation = ( + inspect.signature(litellm.completion) + .parameters["reasoning_effort"] + .annotation + ) + # Optional[Literal[...]] -> unwrap the Union, then read the Literal args. + literals = [ + v + for arg in typing.get_args(annotation) + for v in typing.get_args(arg) + if isinstance(v, str) + ] + return literals or None + except Exception: + return None + + +def _parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]: + """Split ``argv`` into harness options and pass-through script flags.""" + parser = argparse.ArgumentParser( + prog=f"python -m {__spec__.name}" if __spec__ else None, + description=textwrap.dedent(__doc__), + ) + parser.add_argument("script", help="Path to the script to run") + parser.add_argument( + "--model", + type=str, + default=os.environ.get("EFFECTFUL_LLM_MODEL", ""), + help="LLM model to use", + ) + parser.add_argument( + "--num-retries", + type=int, + default=5, + help="Number of retries for malformed/failing LLM output", + ) + parser.add_argument( + "--langfuse", + action="store_true", + help="Whether to log LLM calls and metadata to Langfuse", + ) + parser.add_argument( + "--render", + action="store_true", + help="Live-render the streaming message history in the terminal", + ) + parser.add_argument( + "--dump-system-prompt", + type=str, + default=None, + metavar="PATH", + help="Dump the assembled system prompt to this Markdown file", + ) + parser.add_argument( + "--tool-choice", + type=str, + default="auto", + choices=["required", "auto", "none"], + help="Whether to require, allow, or disable tool calls (none means disabled)", + ) + parser.add_argument( + "--reasoning-effort", + type=str, + default=None, + choices=_reasoning_effort_choices(), + help="Reasoning effort forwarded to litellm.completion", + ) + parser.add_argument( + "--eval-provider", + type=str, + default="unsafe", + choices=sorted(EVAL_PROVIDERS), + help="Provider that runs model-authored Python", + ) + parser.add_argument( + "--pdb", + action="store_true", + help="Drop into pdb post-mortem on an unhandled error (like `python -m pdb`)", + ) + parser.add_argument( + "--persist-db", + type=str, + default=None, + metavar="PATH", + help=( + "Checkpoint persisted Agent state/history to this SQLite database " + "(installs SQLitePersister)" + ), + ) + return parser.parse_known_args(argv) + + +def main(argv: list[str] | None = None) -> None: + ns, script_args = _parse_args(sys.argv[1:] if argv is None else argv) + # The script should see only its own flags, under its own name. + sys.argv = [ns.script, *script_args] + # Mirror `python