diff --git a/.github/workflows/nightly_examples.yml b/.github/workflows/nightly_examples.yml new file mode 100644 index 000000000..294185174 --- /dev/null +++ b/.github/workflows/nightly_examples.yml @@ -0,0 +1,50 @@ +name: Nightly Examples + +# Every example in docs/source/llm_examples, run against a live model with its +# own default arguments. +# +# Pull requests get two cheaper checks: the Test workflow proves every example +# still imports and parses its flags (no API key needed), and the LLM +# Integration Tests workflow runs basics/ live. This is where the rest -- the +# agent searches and multi-round pipelines marked nightly_example -- get run, +# nightly, where a flake costs a re-run rather than a blocked PR. + +on: + schedule: + # 07:00 UTC, after the US evening and before the European morning. + - cron: "0 7 * * *" + workflow_dispatch: + inputs: + model: + description: "Model to run the examples against" + required: false + default: "gpt-4o-mini" + +jobs: + examples: + runs-on: ubuntu-latest + timeout-minutes: 180 + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + python-version: "3.13" + + - name: Set up Python + run: uv python install + + - name: Install Python dependencies + run: uv sync --all-extras --dev + + # No -m filter: the nightly is the run that covers everything. Each + # example is a subprocess, so xdist gets real parallelism across them. + - name: Run every example + env: + EFFECTFUL_LLM_MODEL: ${{ github.event.inputs.model || 'gpt-4o-mini' }} + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + run: | + uv run pytest tests/test_handlers_llm_examples.py -v -n logical 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.yml b/.github/workflows/test.yml index 5743ccbf4..9cdb399de 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -26,7 +26,9 @@ jobs: - name: numpyro paths: tests/test_handlers_numpyro*.py - name: llm - paths: tests/test_handlers_llm*.py + paths: tests/test_handlers_llm*.py -k "not mypy" + - name: llm-mypy + paths: tests/test_handlers_llm*.py -k "mypy" - name: examples paths: tests/test_examples_*.py - name: core 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/.github/workflows/test_llm.yml b/.github/workflows/test_llm.yml index c799de9cf..b88b28aed 100644 --- a/.github/workflows/test_llm.yml +++ b/.github/workflows/test_llm.yml @@ -12,9 +12,12 @@ jobs: test-llm: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: python-version: ["3.13", "3.14"] model: ["gpt-4o-mini"] + mypy: ["mypy", "not mypy"] + name: "test-llm (py ${{ matrix.python-version }}, ${{ matrix.mypy }})" steps: - uses: actions/checkout@v4 @@ -37,4 +40,35 @@ jobs: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} run: | - uv run pytest tests/test_handlers_llm_provider.py tests/test_handlers_llm_tool_calling_poem.py tests/test_handlers_llm_tool_calling_book.py tests/test_handlers_llm_encoding.py -v --tb=short + uv run pytest tests/test_handlers_llm_* -v -n logical \ + -m "not example" -k "${{ matrix.mypy }}" + + # The basics/ examples, run live the way a reader runs them. + examples: + runs-on: ubuntu-latest + timeout-minutes: 30 + name: "basics examples (live)" + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + python-version: "3.13" + + - name: Set up Python + run: uv python install + + - name: Install Python dependencies + run: | + uv sync --all-extras --dev + + - name: Run the basics examples + env: + EFFECTFUL_LLM_MODEL: gpt-4o-mini + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + run: | + uv run pytest tests/test_handlers_llm_examples.py -v -n logical \ + -m "not nightly_example" diff --git a/README.rst b/README.rst index 257f8b597..aae5c9dba 100644 --- a/README.rst +++ b/README.rst @@ -131,7 +131,7 @@ Claude Code, Cursor, Copilot) to: https://basisresearch.github.io/effectful/llm.md This file is auto-generated from the `LLM tutorial notebook `_ -and contains complete API usage examples for templates, tool calling, +and contains complete API usage examples for skills, tool calling, structured output, retries, and more. Learn More diff --git a/docs/source/conf.py b/docs/source/conf.py index 7a33f6058..be54e2e36 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -44,6 +44,12 @@ # Enable documentation inheritance autodoc_inherit_docstrings = True +# Render the ``__init__`` docstring alongside the class docstring. The LLM +# harness handlers rely on this split: a handler's class docstring is injected +# into the model's system prompt, so constructor documentation -- which only +# the caller of ``Handler(...)`` can act on -- lives on ``__init__`` instead. +autoclass_content = "both" + # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # diff --git a/docs/source/effectful.rst b/docs/source/effectful.rst index f6588ea2a..17c97d4ea 100644 --- a/docs/source/effectful.rst +++ b/docs/source/effectful.rst @@ -47,35 +47,172 @@ LLM :members: :undoc-members: -Template -"""""""" +Types +""""" -.. automodule:: effectful.handlers.llm.template +.. automodule:: effectful.handlers.llm.types :members: :undoc-members: -Encoding -"""""""" +Harness +""""""" -.. automodule:: effectful.handlers.llm.encoding +.. automodule:: effectful.handlers.llm.harness :members: :undoc-members: -Completions -""""""""""" - -.. automodule:: effectful.handlers.llm.completions +Command-line launcher +~~~~~~~~~~~~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.__main__ :members: :undoc-members: -Evaluation -"""""""""" - -.. automodule:: effectful.handlers.llm.evaluation +Hooks +~~~~~ + +.. automodule:: effectful.handlers.llm.harness.hooks :members: :undoc-members: - +Serialization +~~~~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.serialization + :members: + :undoc-members: + +Provision +~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.provision + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.provision.litellm + :members: + :undoc-members: + +Legibility +~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.legibility + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.legibility.framework + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.legibility.lexical + :members: + :undoc-members: + +Execution +~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.execution + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.execution.hooks + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.execution.builtin + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.execution.restricted + :members: + :undoc-members: + :private-members: + +Validation +~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.validation + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.validation.hooks + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.validation.pydantic + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.validation.mypy + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.validation.ty + :members: + :undoc-members: + +Synthesis +~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.synthesis + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.synthesis.snippet + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.synthesis.function + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.synthesis.body + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.synthesis.toolcall + :members: + :undoc-members: + +Durability +~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.durability + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.durability.transaction + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.durability.retrying + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.durability.persistence + :members: + :undoc-members: + +Observability +~~~~~~~~~~~~~ + +.. automodule:: effectful.handlers.llm.harness.observability + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.observability.rich + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.observability.dump + :members: + :undoc-members: + +.. automodule:: effectful.handlers.llm.harness.observability.langfuse + :members: + :undoc-members: + + Jax ^^^ 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/acp/__init__.py b/docs/source/llm_examples/acp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/llm_examples/acp/assistant.py b/docs/source/llm_examples/acp/assistant.py new file mode 100644 index 000000000..cbb018349 --- /dev/null +++ b/docs/source/llm_examples/acp/assistant.py @@ -0,0 +1,93 @@ +"""A coding assistant served over the Agent Client Protocol. + +Demonstrates: +- An ``Agent`` exposed to any ACP editor (Zed, VS Code, Obsidian, Emacs) +- Editor capabilities offered to the model as ordinary ``Tool``\\ s +- Streaming model output and live tool-call status as ``session/update`` +- A prompt as typed parts -- prose, ``Attachment`` references, images -- rather + than one flattened string, so an attached file costs a line of context and a + screenshot arrives as an image the model actually sees + +All the protocol machinery lives in ``library.py``; what is left here is an agent, +four imported tools, and a command line -- which is the point of the example. + +The ``prompt`` skill's signature is the server's contract (see +`EffectfulACPAgent`): the name matches the protocol method it answers +(``session/prompt``), and the parameters are what an editor's prompt can carry. +Attachments arrive by reference -- the model reads one through the editor if the +request needs it -- which is what keeps a large attached file from being frozen +into the conversation whole. + +``acp_ask_user`` is one of the four tools, and the only one that is a question +rather than an action: it needs the client's ``elicitation.form`` capability, and +reports itself as a failed call in an editor that has none, so importing it costs +nothing where it will not work. + +The imports are what put the editor's capabilities in the assistant's lexical scope, +and lexical scope is how a `Skill` finds its tools, so they are load-bearing despite +looking unused. `EffectfulACPAgent` is imported inside `main` instead, to keep the +protocol out of the scope the model is shown. + +This is a server: it speaks JSON-RPC over stdin and stdout and expects an editor at +the other end, so it is not useful to run from a terminal by hand. Configure the +editor to launch it, with the command line it would run being:: + + python -m effectful.handlers.llm.harness docs/source/llm_examples/acp/assistant.py \\ + --persist-db /tmp/acp_sessions.db + +``--persist-db`` is what makes a session survive a restart: the editor's session id +becomes the agent's ``__agent_id__``, and the conversation is checkpointed under it +and replayed when the editor reopens the session. + +Set ``ACP_OFFER_MODELS`` to a comma-separated list to put a picker in the editor's +UI, so the session can be switched without editing the editor's configuration. +""" + +import argparse +import asyncio +import collections.abc +import dataclasses + +import PIL.Image +from library import ( # noqa: F401 + Attachment, + acp_ask_user, + acp_read_text_file, + # acp_run_terminal_command, + acp_update_plan, + acp_write_text_file, +) + +from effectful.handlers.llm import Skill + + +@dataclasses.dataclass +class Assistant: + """A coding assistant working inside the user's editor. + + Your reply is shown to the user in their editor and rendered as Markdown, so + write it as prose addressed to them; headings, lists, links and fenced code + blocks all display. + """ + + __agent_id__: str = "" + + @Skill.define + def prompt( + self, + user_input: str, + attachments: collections.abc.Sequence[Attachment] = (), + images: collections.abc.Sequence[PIL.Image.Image] = (), + ) -> str: + """{user_input}{attachments}{images}""" + + +def main() -> None: + from library import EffectfulACPAgent + + argparse.ArgumentParser(description=__doc__).parse_args() + asyncio.run(EffectfulACPAgent(Assistant).serve()) + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/acp/library.py b/docs/source/llm_examples/acp/library.py new file mode 100644 index 000000000..682556715 --- /dev/null +++ b/docs/source/llm_examples/acp/library.py @@ -0,0 +1,2805 @@ +"""Serve an `effectful.handlers.llm` `Agent` over the Agent Client Protocol. + +The [Agent Client Protocol](https://agentclientprotocol.com) (ACP) is how an editor +-- Zed, VS Code, Obsidian, Emacs -- drives a coding agent: a JSON-RPC conversation +over stdio in which the editor opens a session, sends prompts, and receives a stream +of updates describing what the agent is doing. This module makes an `Agent` speak it. + +## The shape of the thing + +An ACP server has to answer three questions that the harness already answers as +effects, so almost nothing here is new machinery -- it is three translations: + +| ACP concept | effectful concept | Where | +| ----------- | ----------------- | ----- | +| `session/update` notifications | the `completion` and `call_tool` effects | `ACPSessionReporter` | +| `session/request_permission` | intercepting `call_tool` | `ACPPermissionGate` | +| `fs/*`, `terminal/*` and `elicitation/*` client methods | `Tool`s the model may call | `ACPToolRuntime` | + +`ACPSessionReporter` is the interesting one, and it is a close sibling of +`~effectful.handlers.llm.harness.observability.rich.RichTerminalRenderer`: both force +`completion` onto the streaming path and re-render the deltas somewhere. One renders +to a terminal, this one to a JSON-RPC pipe. + +## Threads + +The harness is synchronous -- `litellm.completion` blocks, and the completion loop is +a `while` -- while ACP is asyncio. So a prompt runs in a worker thread +(`asyncio.to_thread`), and `ACPSession` is the only thing that crosses back: + +- **Notifications** are enqueued (`ACPSession.notify`) and drained by one writer task + per session, so they cannot interleave on the wire. +- **Requests** block the worker until the editor answers (`ACPSession.call`). + +`asyncio.to_thread` copies the caller's `contextvars`, and the handler stack lives in +one (`effectful.internals.runtime.INTERPRETATION`), so the worker inherits whatever +stack was installed around the server and adds this session's handlers on top. + +## What the protocol asks for that an effect does not supply + +The three translations are most of it, but not all: ACP also has requirements about +the *conversation*, which `EffectfulACPAgent` is where it meets. The ones with teeth, +each of which has a test: + +- A session is opened on a directory (`cwd`), and everything in it -- what the model + is told it is working on, where a terminal command runs -- is rooted there rather + than in whatever directory the editor happened to spawn this process from. +- Every prompt is answered with a `StopReason`, and the interesting ones are not + `end_turn`: a reply cut off at the token limit, one the provider refused, one the + user cancelled. +- Whatever the turn told the editor, it told it *before* answering; and no tool call + it announced is left without a terminal status. +- A capability is claimed if and only if it is implemented, in both directions -- + including `session/list`, which is claimed only when there is somewhere to keep + the answer, since a conversation the agent cannot name is one it cannot reopen. +""" + +import asyncio +import base64 +import collections.abc +import concurrent.futures +import contextlib +import dataclasses +import datetime +import functools +import io +import json +import os +import sqlite3 +import sys +import threading +import typing +import urllib.parse +import urllib.request +import uuid + +import acp +import acp.interfaces +import acp.schema +import litellm +import pydantic +import pydantic_core +from PIL import Image + +from effectful.handlers.llm import Agent, Encodable, Tool +from effectful.handlers.llm.harness.durability.persistence import SQLitePersister +from effectful.handlers.llm.harness.hooks import ( + PromptInjectingInterpretation, + ToolCallExecutionError, + call_system, + call_tool, + completion, +) +from effectful.handlers.llm.harness.serialization import ( + DecodedToolCall, + PromptSection, + to_content_blocks, +) +from effectful.handlers.llm.harness.synthesis.body import FinalBodySynthesizer +from effectful.handlers.llm.harness.synthesis.snippet import StatefulReplSynthesizer +from effectful.ops.semantics import coproduct, fwd, handler +from effectful.ops.syntax import ObjectInterpretation, implements +from effectful.ops.types import Interpretation + +type ContentBlock = ( + acp.schema.TextContentBlock + | acp.schema.ImageContentBlock + | acp.schema.AudioContentBlock + | acp.schema.ResourceContentBlock + | acp.schema.EmbeddedResourceContentBlock +) + +type ToolCallContent = ( + acp.schema.ContentToolCallContent + | acp.schema.FileEditToolCallContent + | acp.schema.TerminalToolCallContent +) + +type ConfigOption = ( + acp.schema.SessionConfigOptionSelect | acp.schema.SessionConfigOptionBoolean +) + +type ElicitationProperty = ( + acp.schema.ElicitationStringPropertySchema + | acp.schema.ElicitationBooleanPropertySchema + | acp.schema.ElicitationMultiSelectPropertySchema +) + +# --------------------------------------------------------------------------- +# The editor's own capabilities, offered to the model as tools +# +# Declared here and given meaning by `ACPToolRuntime`, one instance per session. +# A `Tool` is an `Operation`, so the declaration is the signature and the +# docstring -- the two things the model is shown -- and the body is unreachable +# under any interpretation that handles it. Splitting them that way is what lets +# the tools be plain module-level names: a script imports them to put them in its +# skills' lexical scope, and the session supplies the editor they talk to. +# --------------------------------------------------------------------------- + + +READ_RESULT_LIMIT_LINES = 1000 +"""How many lines a read may return before it is truncated with a notice. + +The bound that keeps one tool call from swallowing a conversation's budget: a +read lands in the history and is re-sent with every request after, so an +unbounded read of a large file is a large recurring cost incurred in one call. +The notice tells the model how to page with ``line``/``limit``; a model that +passes its own ``limit`` has chosen a size, and is not second-guessed. +""" + + +@Tool.define +def acp_read_text_file( + path: str, line: int | None = None, limit: int | None = None +) -> str: + """Read a text file and return its contents. + + Reads through the user's editor, so it sees unsaved changes in an open buffer + -- use it in preference to opening the file yourself. `path` must be + absolute. `line` is a 1-based line to start from, and `limit` a maximum + number of lines to return; pass neither to read from the start. A long + read with no `limit` is truncated after 1000 lines, with a notice saying + how to read on from where it stopped. + """ + raise RuntimeError("missing handler") + + +@Tool.define +def acp_write_text_file(path: str, content: str) -> str: + """Write `content` to the file at absolute `path`, replacing what is there. + + Writes through the user's editor, so the change lands in the open buffer and + the user sees it as an ordinary edit they can undo. Read the file first + unless you are creating it. + """ + raise RuntimeError("missing handler") + + +@Tool.define +def acp_run_terminal_command(command: str, args: list[str]) -> str: + """Run `command` with `args` in a terminal and return its output. + + The terminal is the user's own, so they can watch the command run. Blocks + until it exits. Prefer this to a shell one-liner assembled by hand: `args` + are passed to the process directly, with no shell to quote for. + """ + raise RuntimeError("missing handler") + + +@pydantic.dataclasses.dataclass +class PlanStep: + """One step of the plan you are showing the user. + + Deliberately not `acp.schema.PlanEntry`, which is the same three fields plus the + ``_meta`` every ACP type carries. A tool's parameters are turned into a strict JSON + Schema, and strict schemas require every property, so the model would be obliged to + supply a value for a field whose own documentation says implementations must not + assume anything about it. The two vocabularies below are borrowed from the protocol + rather than restated, so the part that could drift cannot. + """ + + content: str + """What you will do, in a few words, as you would say it to them.""" + + priority: acp.schema.PlanEntryPriority = "medium" + status: acp.schema.PlanEntryStatus = "pending" + + +@Tool.define +def acp_update_plan(steps: list[PlanStep]) -> str: + """Show the user your plan for a job with several steps, and keep it current. + + The editor renders this as a checklist beside the conversation, so the user can + see where you are without reading back through it. Send the *whole* plan every + time -- each call replaces the last -- marking at most one step `in_progress` and + everything you have finished `completed`. + + Worth doing for work that takes several tool calls, and not worth it otherwise: a + one-step plan tells the user nothing they did not just ask for. + """ + raise RuntimeError("missing handler") + + +@pydantic.dataclasses.dataclass +class AskField: + """One thing you are asking the user for, rendered as a field in a form. + + ACP describes a form as a JSON Schema, and `_elicitation_property` is where this + becomes one. Handing the model that schema directly instead -- an + `acp.schema.ElicitationSchema`, whose ``properties`` is a dict of a seven-way union + of property types -- would ask it to author JSON Schema as a side errand of asking + a question, and costs an order of magnitude more of its context to describe. This + is the small closed subset an editor can actually render. + """ + + name: str + """The key its answer comes back under. Short, and unique within one question.""" + + title: str + """The label the editor shows beside it, as you would say it to them.""" + + description: str = "" + """Optional detail, for a label that needs more than a few words.""" + + kind: typing.Literal["text", "choice", "boolean"] = "text" + """What kind of answer: free text, one of `choices`, or a yes/no.""" + + choices: list[str] = dataclasses.field(default_factory=list) + """The options, when `kind` is ``choice``. Required for one, ignored otherwise.""" + + required: bool = True + """Whether the user must fill it in before the form can be submitted.""" + + +@Tool.define +def acp_ask_user(message: str, fields: list[AskField]) -> str: + """Ask the user something, as a small form in their editor, and wait for a reply. + + For a decision that is genuinely theirs to make: two defensible ways to do what + they asked, a destructive change worth confirming, a preference you cannot read + off the code. Say in `message` what you are about to do and why the answer + matters, and keep `fields` to the one or two things you actually need. + + Not for anything you could find out yourself. If the answer is in a file, read + the file -- asking instead spends the user's attention to save you a tool call, + and an assistant that asks before every step is worse than one that gets on with + it and says what it assumed. + + The user may decline, and declining is an answer: you are told so, and should + then continue without it or explain what you cannot decide. They may also dismiss + the form, which ends the turn. + """ + raise RuntimeError("missing handler") + + +# --------------------------------------------------------------------------- +# The knobs a session offers the user +# +# Two protocol features that cost an agent almost nothing and that an editor +# renders for free, both of which stay dead until the agent describes itself: +# a *mode* (`session/set_mode`) and a *config option* (`session/set_config_option`). +# Only `select` options are used, deliberately -- a boolean one is gated behind +# the client's `session.configOptions.boolean` capability, and most clients, +# including VS Code's, advertise no session capabilities at all. +# --------------------------------------------------------------------------- + +ASK, AUTO, PLAN = "ask", "auto", "plan" + +SESSION_MODES: tuple[acp.schema.SessionMode, ...] = ( + acp.schema.SessionMode( + id=ASK, + name="Ask", + description="Ask before running each tool.", + ), + acp.schema.SessionMode( + id=AUTO, + name="Auto", + description="Run tools without asking. Undo is your editor's.", + ), + acp.schema.SessionMode( + id=PLAN, + name="Plan", + description="Read and discuss, but change nothing: no writes, no commands.", + ), +) + +MUTATING_TOOLS = frozenset( + {acp_write_text_file.__name__, acp_run_terminal_command.__name__} +) +"""What `PLAN` mode refuses: the tools that change the user's editor or machine. + +A denylist of the two `ACPToolRuntime` offers, and deliberately not a sandbox. The +harness may also be running model-authored Python -- `exec_code`, +``write_and_run_body`` -- which this says nothing about, because that is the eval +provider's business and the launcher's ``--eval-provider none`` is the switch for it. +Naming the mode "Plan" rather than "read-only" is what keeps that promise honest. +""" + +UNGATED_TOOLS = frozenset({acp_ask_user.__name__}) +"""What `ACPPermissionGate` lets through without asking: asking the user something. + +The gate exists to put a question in front of the user before a tool runs, so +prompting for permission to ask them a question is the one case where it defeats +itself -- a dialog about a dialog, answered by the same person, immediately followed +by the real one. + +Safe on its own terms rather than by exception. The tool's only effect is a form on +the user's screen: it reads nothing, changes nothing, and dismissing it already +cancels the turn, so the decision the gate would have offered is one the user still +has. `PLAN` mode leaves it alone for the same reason -- asking changes nothing, and a +session that may not act is exactly where a clarifying question is worth most. +""" + +MODE_OPTION_ID = "mode" +MODEL_OPTION_ID = "model" +INHERIT_MODEL = "" +"""The `model` option's value meaning "whatever the process was configured with". + +An empty string rather than the model's name, because this agent does not know that +name: the model is bound into `LiteLLMConfigurer` by whoever assembled the stack, and +nothing in the protocol layer can see it. Saying "as configured" is honest; naming a +model here would be a guess printed in the user's editor. +""" + +OFFER_MODELS_ENV = "ACP_OFFER_MODELS" +"""Environment variable naming the models the editor's picker should offer. + +An environment variable rather than a flag because an editor launches an agent with a +command and an environment, and the model this one *starts* on already comes from the +environment -- the launcher's ``--model`` defaults to ``EFFECTFUL_LLM_MODEL``. Putting +the models it can switch to anywhere else would split one setting across two +mechanisms in the same block of the editor's configuration. +""" + + +def _offered_models() -> tuple[str, ...]: + """The picker's models as named in the environment, in order, or none. + + Comma-separated, since a model name may contain ``/``, ``-``, ``.`` and ``:`` but + never a comma. Blanks are dropped rather than offered: a picker whose list has an + empty entry in it is worse than no picker, and a trailing comma is the likeliest + way to write one by accident. + """ + listed = os.environ.get(OFFER_MODELS_ENV, "").split(",") + return tuple(model.strip() for model in listed if model.strip()) + + +FLUSH_TIMEOUT = 5.0 +"""How long a turn waits for its queued updates to reach the editor before answering. + +A courtesy wait, not a correctness one: ACP requires the updates to be *sent* before +the final response, and this is what makes that a wait rather than a hope. But an +editor that has stopped reading its own pipe should not be able to wedge the turn +trying to tell it so, which is the whole reason for the bound. Nothing a user waits +on is measured by it, so it is short and fixed rather than configurable. +""" + + +# --------------------------------------------------------------------------- +# Per-session state, and the crossing between the event loop and the harness +# --------------------------------------------------------------------------- + + +class SessionCancelled(BaseException): + """Raised in the worker thread when the editor sends ``session/cancel``. + + Raising is how a cancellation noticed deep in the completion loop -- inside a + tool, between stream chunks, while waiting on the editor -- reaches + `EffectfulACPAgent.prompt`, which is the only place that may answer the request. + + Deriving from `BaseException` rather than `Exception` is load-bearing. + `~effectful.handlers.llm.harness.durability.retrying.TenacityRetryer` catches + `Exception`-derived tool failures and hands them to the model as feedback, and + retries `Exception`-derived completion failures, so a cancellation raised inside + the loop would otherwise be swallowed and reported to the model as a broken tool + rather than stopping the turn. + """ + + +@dataclasses.dataclass +class ACPSession[A: Agent]: + """Everything the server keeps for one ACP session. + + ACP has no such object -- it addresses sessions by id and leaves the rest to the + agent -- so this is where the per-session state lives: + + * the `Agent` whose history *is* the conversation, and whose ``__agent_id__`` is + the session id the editor knows it by; + * the binding to the editor (`client`, `client_capabilities`), which every call + back to it needs; + * the directories the session is *about* (`cwd`, `additional_directories`), which + ACP calls its root set; + * the lock that keeps two prompts on one session from running at once (`lock`), + since they would put two worker threads on one agent's history; + * the queue of pending notifications and the task draining it (`updates`, + `notify`, `writer`, `drain`). + + Most of the methods serve the second, because the harness is synchronous and the + protocol is not: a prompt runs in a worker thread while the connection lives on + the event loop, so every call back to the editor crosses a thread boundary, and + that crossing is written here once. + + Construct it on the event loop thread: it captures the running loop and starts + the writer task on it. + """ + + agent: A + client: acp.interfaces.Client + client_capabilities: acp.schema.ClientCapabilities + + cwd: str = "" + """The session's working directory, as an absolute path. + + ACP requires it of every ``session/new`` and ``session/load``: it "MUST be used + for the session regardless of where the Agent subprocess was spawned", and is the + base every relative path in the session resolves against. The agent's own process + directory is whatever the editor happened to be launched from and is never it, + which is why nothing here consults `os.getcwd`. + """ + + additional_directories: tuple[str, ...] = () + """Further absolute paths the session may work in, beyond `cwd`.""" + + mode_id: str = ASK + """Which of `SESSION_MODES` the user has picked. Read by `ACPPermissionGate`.""" + + model: str = INHERIT_MODEL + """The model the user picked, or `INHERIT_MODEL` for the configured one.""" + + title: str = "" + """A human-readable name for the conversation, taken from its first prompt.""" + + poll_interval: float = 0.1 + """How often a worker thread waiting on the editor re-reads `cancel`.""" + + loop: asyncio.AbstractEventLoop = dataclasses.field( + default_factory=asyncio.get_running_loop + ) + cancel: threading.Event = dataclasses.field(default_factory=threading.Event) + updates: asyncio.Queue = dataclasses.field(default_factory=asyncio.Queue) + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + writer: asyncio.Task | None = None + + def __post_init__(self) -> None: + """Start the one task that drains `updates`. + + It belongs to the session's whole life rather than to a turn: `load_session` + replays a conversation with no turn in progress, and a turn's own last act is + to wait for the queue to empty. Starting it per turn would make the first of + those deliver nothing and the second wait on a queue nobody is reading. + """ + self.writer = self.loop.create_task(self.drain()) + + @property + def session_id(self) -> str: + """The ACP session id this object represents.""" + return self.agent.__agent_id__ + + @property + def fs_capabilities(self) -> acp.schema.FileSystemCapabilities: + """What the editor will do to files on the agent's behalf. + + ``fs`` is optional in `acp.schema.ClientCapabilities`, and a client that sends + it as ``null`` means the same thing as one that claims nothing: reading it + through here answers "no" rather than raising `AttributeError` inside whatever + tool happened to ask. + """ + return self.client_capabilities.fs or acp.schema.FileSystemCapabilities() + + @property + def elicitation_capabilities(self) -> acp.schema.ElicitationCapabilities: + """Which kinds of question the editor will put to the user for us. + + Optional exactly as `fs` is, and read through here for the same reason: a + client that sends it as ``null`` means the same thing as one that claims + nothing, and `acp_ask_user` should be told "no" rather than meet an + `AttributeError`. + + There is no agent-side capability to match this one, so nothing in + `EffectfulACPAgent.agent_capabilities` changes: elicitation is something the + *client* offers, and the claim-iff-implemented rule applies to it only + inbound -- ask before asking, and take no for an answer. + """ + return ( + self.client_capabilities.elicitation or acp.schema.ElicitationCapabilities() + ) + + @property + def roots(self) -> tuple[str, ...]: + """Every directory this session may work in, `cwd` first. + + ACP calls this the session's effective root set, and requires `cwd` to be part + of it -- so it is derived here rather than stored, and cannot fall out of step + with `cwd`. + """ + return (self.cwd, *self.additional_directories) if self.cwd else () + + @functools.cached_property + def reporter(self) -> "ACPSessionReporter": + """The reporter in `intp`, reachable by name. + + It is the only handler of the three that outlives the call it is installed + for: it counts a turn's requests and remembers which tool calls it has + announced, and `EffectfulACPAgent.prompt` reads both once the worker is done. + """ + return ACPSessionReporter(self) + + @functools.cached_property + def intp(self) -> Interpretation: + """The handler stack that runs a prompt on this session. + + Built once per session, and installed on top of whatever stack the process is + already running under (see `EffectfulACPAgent._answer`). Every handler in it + closes over this session, which is how three module-level concerns -- the + tools, the reporting, the permission gate -- reach *this* editor without any + of them being per-session types. + + `ACPPermissionGate` goes on last, so it is outermost: it must decide about a + call before `ACPSessionReporter` announces it as running. + """ + h = coproduct( + self.reporter, + ACPToolRuntime(self), + ) + h = coproduct(h, ACPSessionConfig(self)) + h = coproduct(h, ACPPermissionGate(self)) + return h + + def notify(self, update: typing.Any) -> None: + """Enqueue a `session/update` for the writer task. Never blocks. + + Ordering is the reason for the queue. Scheduling each notification as its own + coroutine would let two of them interleave inside the connection's writer; + one consumer draining a queue keeps them in the order they were produced. + + The enqueue is immediate when this is already the loop's own thread, and only + deferred when it is not. `call_soon_threadsafe` unconditionally would defer it + past the caller's own next await, so a caller that notifies and then waits for + the queue to empty -- `load_session` does exactly that -- would find an empty + queue and return before anything had been put on it. + """ + try: + on_loop = asyncio.get_running_loop() is self.loop + except RuntimeError: # no running loop: definitely a worker thread + on_loop = False + if on_loop: + self.updates.put_nowait(update) + else: + self.loop.call_soon_threadsafe(self.updates.put_nowait, update) + + def call[T]( + self, + coro: typing.Coroutine[typing.Any, typing.Any, T], + *, + orphan: collections.abc.Callable[[concurrent.futures.Future], None] + | None = None, + ) -> T: + """Run a client request on the event loop and wait for the editor's answer. + + Used for the requests whose *answer* the worker needs -- a permission + decision, a file's contents -- as opposed to the notifications above. + + The wait is unbounded and interruptible, which is the pairing the protocol + asks for. Unbounded because the thing most often waited on here is a human: + ``session/request_permission`` puts a dialog in front of the user and there is + no honest deadline for reading it. A bound would eventually tell the model the + editor never answered while the dialog was still on screen, and then refuse + the call out from under the user about to approve it. Every other ACP agent + simply waits, and so does this one. + + Interruptible is what makes waiting forever safe, and it is not optional. A + plain ``.result()`` parks the worker thread with no way back: that thread holds + the session lock, so every later prompt blocks behind it, and the turn never + reaches the points where it reads `cancel` -- before a completion, before a + tool call, between stream chunks. `session/cancel` would be a lie in exactly + the case a user most wants it, a prompt they cannot or will not answer. + Waiting in short slices and re-reading the flag is what makes it true, and it + leaves the user, rather than a clock, deciding when a silent editor has waited + long enough. + + Cancellation normally also cancels the request itself -- the polite thing + for a permission dialog the user has just walked away from. ``orphan`` is + for the requests where that would *lose* something: a cancelled + ``terminal/create`` was already sent, the editor allocates the terminal + and answers, and an agent that cancelled the answer has leaked a terminal + it never learned the id of. With ``orphan`` given, cancellation leaves the + request running and attaches the callback to its eventual completion, so + the caller can dispose of whatever the answer turns out to be. + + Raises: + SessionCancelled: If the turn was cancelled while waiting. + """ + future = asyncio.run_coroutine_threadsafe(coro, self.loop) + while True: + if self.cancel.is_set(): + if orphan is None: + future.cancel() + else: + future.add_done_callback(orphan) + raise SessionCancelled + with contextlib.suppress(concurrent.futures.TimeoutError): + return future.result(timeout=self.poll_interval) + + def detach( + self, coro: typing.Coroutine[typing.Any, typing.Any, typing.Any] + ) -> None: + """Send a client request without waiting for -- or ever cancelling -- it. + + For requests that are obligations rather than questions: the answer is + not needed and the sending must not depend on the turn's fate. Releasing + a terminal is the canonical case -- it belongs to *whichever* way the + turn ends, so tying it to an interruptible wait would let the very + cancellation that ends a command also revoke the release it owes. + """ + with contextlib.suppress(RuntimeError): # a loop already shut down + asyncio.run_coroutine_threadsafe(coro, self.loop) + + async def flush(self) -> None: + """Wait until every notification produced so far has reached the editor. + + ACP requires this of both requests that report progress: an agent "MAY send + update notifications before responding, but MUST do so before the final + response", and `session/load` must finish streaming the conversation before it + answers. The queue and its writer are what make that a wait rather than a + guarantee, so the wait is written once and used by both. + + Bounded by `FLUSH_TIMEOUT`, because it is a *courtesy* wait: the alternative to + answering a little early is never answering at all, and an editor that has + stopped reading its own pipe should not be able to wedge the turn that is + trying to tell it so. Unlike `call`, nobody is being waited *for* here -- the + updates have already been produced -- so a bound costs the user nothing. + """ + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self.updates.join(), timeout=FLUSH_TIMEOUT) + + async def drain(self) -> None: + """Deliver queued notifications until cancelled. One task per session. + + ``task_done`` is what lets ``updates.join()`` return, and a turn ends by + awaiting exactly that, so it has to happen even for a notification that failed + to send -- otherwise one dropped update would deadlock the end of that turn + and every turn after it. + """ + while True: + update = await self.updates.get() + try: + await self.client.session_update(self.session_id, update) + except Exception: + pass + finally: + self.updates.task_done() + + +def _elicitation_property(field: AskField) -> ElicitationProperty: + """One `AskField` as the JSON Schema property an editor renders a widget from. + + ACP restricts these to primitives, so this is a small closed mapping rather than + a general schema translation: text, a single choice, a yes/no. `oneOf` carries the + choices rather than `enum` because its options are titled, and a title is what the + editor puts on the control. + + A ``choice`` with nothing to choose from would render as an empty dropdown, so it + degrades to a text field: the model asked for an answer, and a control the user + cannot use would be a worse way to fail than one they can type into. + """ + description = field.description or None + if field.kind == "boolean": + return acp.schema.ElicitationBooleanPropertySchema( + type="boolean", title=field.title, description=description + ) + if field.kind == "choice" and field.choices: + return acp.schema.ElicitationStringPropertySchema( + type="string", + title=field.title, + description=description, + one_of=[ + acp.schema.EnumOption(const=choice, title=choice) + for choice in field.choices + ], + ) + return acp.schema.ElicitationStringPropertySchema( + type="string", title=field.title, description=description + ) + + +def _elicitation_schema(fields: list[AskField]) -> acp.schema.ElicitationSchema: + """The whole form: one property per field, and which of them are required.""" + return acp.schema.ElicitationSchema( + type="object", + properties={field.name: _elicitation_property(field) for field in fields}, + required=[field.name for field in fields if field.required] or None, + ) + + +def _answers_as_text( + fields: list[AskField], content: collections.abc.Mapping[str, typing.Any] | None +) -> str: + """What the user filled in, as lines the model can read. + + Keyed by the fields that were asked rather than by what came back, so the answers + arrive in the order they were asked and a field the user left blank is reported as + blank instead of vanishing -- the model needs to know it asked and got nothing. + + `content` is optional even on an ``accept``, which is the protocol allowing a form + with nothing in it to be submitted. + """ + answers = content or {} + lines = [] + for field in fields: + value = answers.get(field.name) + if isinstance(value, bool): + shown = "yes" if value else "no" + elif isinstance(value, list): + shown = ", ".join(str(item) for item in value) or "(nothing selected)" + elif value is None or value == "": + shown = "(left blank)" + else: + shown = str(value) + lines.append(f"{field.name}: {shown}") + return "The user answered:\n" + "\n".join(lines) + + +@dataclasses.dataclass +class ACPToolRuntime(PromptInjectingInterpretation): + """Your filesystem is the user's editor, not this process's disk. + + The `acp_read_text_file`, `acp_write_text_file` and `acp_run_terminal_command` + tools go through the editor the user is sitting in front of, so a read sees + unsaved changes in an open buffer and a write lands as an edit the user can undo. + Prefer them to anything you might reach for in code. Every path you pass them must + be absolute, and should be inside the directories listed below; a terminal command + runs in the first of those. + + The user is sitting there too, so `acp_ask_user` can put a question to them and + wait for the answer -- worth it for a decision that is theirs, and not for + anything you could learn by reading a file. + """ + + session: ACPSession + + def _directories_section(self) -> PromptSection: + """The session's root set, named. Computed, so it cannot be a docstring. + + This is the whole point of ACP handing `cwd` to `session/new`: without it the + model has no idea which project it is in, and "use absolute paths" is advice + it cannot act on. + """ + roots = self.session.roots + listed = ( + "\n".join(f"- `{root}`" for root in roots) + if roots + else "- (the editor named none; ask the user before assuming a path)" + ) + return PromptSection( + type="prompt_section", + title="Directories this session is about", + content=to_content_blocks( + "The user opened this session on the following directories. The first " + "is the working directory: it is what a relative path would mean, and " + "it is where a terminal command runs.\n\n" + listed + ), + ) + + @implements(call_system) + def call_system( + self, harness_prompt: PromptSection, agent_prompt: PromptSection + ) -> typing.Any: + """Add the session's directories, then the class docstring the base adds. + + Appending before delegating puts them immediately ahead of the docstring's + section, so the sentence there about "the directories listed below" is + followed by the list. + """ + return super().call_system( + PromptSection( + type="prompt_section", + title=harness_prompt["title"], + content=[*harness_prompt["content"], self._directories_section()], + ), + agent_prompt, + ) + + @implements(acp_read_text_file) + def acp_read_text_file( + self, path: str, line: int | None = None, limit: int | None = None + ) -> str: + """Read through the editor, truncating an unbounded read of a long file. + + Only a read the model left unbounded is truncated: an explicit `limit` + already asked the editor for a bounded answer, and cutting it further + would make the parameter mean less than it says. The notice names the + line to continue from, so paging costs the model no arithmetic. + """ + if not self.session.fs_capabilities.read_text_file: + raise NotImplementedError( + "this editor cannot read files on your behalf; ask the user instead" + ) + content = self.session.call( + self.session.client.read_text_file( + self.session.session_id, path, line=line, limit=limit + ) + ).content + if limit is not None: + return content + lines = content.splitlines(keepends=True) + if len(lines) <= READ_RESULT_LIMIT_LINES: + return content + start = line or 1 + shown = start + READ_RESULT_LIMIT_LINES - 1 + total = start - 1 + len(lines) + return "".join(lines[:READ_RESULT_LIMIT_LINES]) + ( + f"\n[truncated: showing lines {start}..{shown} of {total}; call again " + f"with line={shown + 1} (and a limit, if you like) for the rest]" + ) + + @implements(acp_write_text_file) + def acp_write_text_file(self, path: str, content: str) -> str: + if not self.session.fs_capabilities.write_text_file: + raise NotImplementedError( + "this editor cannot write files on your behalf; propose the change " + "to the user as text instead" + ) + # Read before writing, so the editor can draw a before-and-after rather than + # just the new text. One extra round trip to the editor, which is local and + # fast, and it is skipped entirely when there is nothing to compare against. + self.session.reporter.show_diff(path, content, self._previous_text(path)) + self.session.call( + self.session.client.write_text_file(self.session.session_id, path, content) + ) + # ACP's own response is empty, but the tool is declared to return `str` and + # the model is shown whatever it returns: `null` reads as a call that did + # not do anything. + return f"wrote {len(content)} characters to {path}" + + def _previous_text(self, path: str) -> str | None: + """What is in the file now, for a diff to be drawn against, or `None`. + + `None` covers every reason this can fail to be an answer -- the file is being + created, the editor will not read on our behalf, the read failed -- because + none of them is a reason to fail the *write*. A cancellation is not one of + those: `SessionCancelled` derives from `BaseException` and passes through, so a + cancelled turn still stops here. + """ + if not self.session.fs_capabilities.read_text_file: + return None + try: + return self.session.call( + self.session.client.read_text_file(self.session.session_id, path) + ).content + except Exception: + return None + + @implements(acp_update_plan) + def acp_update_plan(self, steps: list[PlanStep]) -> str: + """Replace the plan the editor is showing, and tell the model what it shows. + + Needs no capability check: the client's `plan` capability gates the incremental + `plan_update` and `plan_removed` notifications, not this one, which is why an + editor that advertises nothing still renders it. + """ + self.session.notify( + acp.update_plan( + acp.plan_entry(step.content, priority=step.priority, status=step.status) + for step in steps + ) + ) + done = sum(step.status == "completed" for step in steps) + return f"Showing the user {len(steps)} step(s), {done} of them completed." + + @implements(acp_ask_user) + def acp_ask_user(self, message: str, fields: list[AskField]) -> str: + """Put a form in front of the user and hand their answers back to the model. + + The wait is the point, and it is unbounded: see `ACPSession.call`, whose + argument for a permission dialog is this one word for word. A form is a + question for a person, and a person is allowed to think. + + Three answers come back, and the difference between the last two is the whole + of why this is not just a permission prompt. ``accept`` is the answers. + ``decline`` is the user saying they will not answer *this*, which is a fact + the model should carry on from rather than a broken tool -- so it returns + normally, with prose saying so. ``cancel`` is the form dismissed, which is + the same gesture as dismissing a permission prompt and ends the turn. + + Raises: + NotImplementedError: If the editor renders no forms; reported to the + model as this call's result, like any other missing capability. + SessionCancelled: If the user dismissed the form instead of answering it. + """ + if not self.session.elicitation_capabilities.form: + raise NotImplementedError( + "this editor cannot show the user a form; ask your question in your " + "reply instead, and stop there so they can answer it" + ) + if not fields: + raise ValueError( + "ask for at least one thing; a form with no fields is a dialog the " + "user can only dismiss" + ) + response = self.session.call( + self.session.client.create_elicitation( + message, + acp.schema.ElicitationFormSessionMode( + session_id=self.session.session_id, + # So the editor draws the form in the tool-call row it belongs to, + # rather than as a dialog with no visible cause. `None` outside a + # tool call, which the field allows. + tool_call_id=self.session.reporter.running, + requested_schema=_elicitation_schema(fields), + ), + ) + ) + if response.action == "cancel": + raise SessionCancelled + if response.action == "decline": + return ( + "The user declined to answer. Do not ask again; either continue " + "without their answer, saying what you assumed, or explain what you " + "cannot decide for them." + ) + if response.action != "accept": + # `OtherElicitationResponse` exists for actions added after this was + # written. Reporting the word rather than guessing which of the three it + # resembles is the only honest thing to do with one. + return f"The editor answered with an action this agent does not know: {response.action!r}." + return _answers_as_text(fields, getattr(response, "content", None)) + + @implements(acp_run_terminal_command) + def acp_run_terminal_command(self, command: str, args: list[str]) -> str: + """Run a command in the user's terminal and return its output and status.""" + if not self.session.client_capabilities.terminal: + raise NotImplementedError( + "this editor cannot run terminal commands on your behalf" + ) + + def release_orphan(created: concurrent.futures.Future) -> None: + # A cancellation that lands *during* `terminal/create` interrupts the + # wait below before the terminal id is ever known -- but the request + # was already sent, so the editor allocates a terminal and answers. + # Without this, that answer is discarded and the terminal leaks: the + # `finally` cannot release an id the agent never learned. (Exactly + # this interleaving is routine on a loaded runner, where the user's + # cancel overtakes a worker that has only just announced the call.) + if not created.cancelled() and created.exception() is None: + self.session.detach( + self.session.client.release_terminal( + self.session.session_id, created.result().terminal_id + ) + ) + + terminal = self.session.call( + self.session.client.create_terminal( + self.session.session_id, + command, + args=args, + # Without this the command runs wherever the editor happened to spawn + # this process, which is not the project the user opened. + cwd=self.session.cwd or None, + ), + orphan=release_orphan, + ).terminal_id + # Hand the terminal to the editor before waiting on it: this is the one thing + # in the protocol that renders *while* it happens. The alternative -- what this + # did until now -- is a spinning row for however long the command takes, and + # then its whole output at once. + self.session.reporter.show_terminal(terminal) + try: + exit_status = self.session.call( + self.session.client.wait_for_terminal_exit( + self.session.session_id, terminal + ) + ) + output = self.session.call( + self.session.client.terminal_output(self.session.session_id, terminal) + ) + finally: + # `detach`, not `call`: the release is owed however the turn ended, + # and must not itself be revoked by the cancellation that ended it + # (nor, on the way out of a *successful* command, may a late cancel + # be allowed to discard the output by raising here). + self.session.detach( + self.session.client.release_terminal(self.session.session_id, terminal) + ) + status = ( + f"exit code {exit_status.exit_code}" + if exit_status.signal is None + else f"killed by {exit_status.signal}" + ) + return f"[{status}]\n{output.output}" + + +@dataclasses.dataclass +class ACPSessionConfig(ObjectInterpretation): + """Apply the session's user-chosen settings to the requests it makes. + + Just the model, for now, and it is one line: `LiteLLMConfigurer` merges its own + configuration *under* whatever the request already carries -- "the merge below + lets a value already in `kwargs` stand" -- so a handler installed above it names + the model by naming it. That is what turns the picker in the editor's UI from a + label into a setting. + + A second `LiteLLMConfigurer` in `ACPSession.intp` would be the obvious way to do + that, and cannot be. It binds its configuration at construction, where `intp` is + built once per session and `model` changes whenever the user touches the picker, + so the choice would freeze at whatever it was when the stack was first built. + Reading `session.model` per request is the whole point, and it is also why + `INHERIT_MODEL` can be honoured at all: `LiteLLMConfigurer` has no way to say + "no opinion" -- its `model` defaults to ``gpt-4o``, which would silently overrule + the model the launcher was configured with. + + Separate from `ACPSessionReporter`, which also handles `completion`, because these + are opposite directions: the reporter watches a request go past and describes it, + this alters it. + """ + + session: ACPSession + + @implements(completion) + def completion(self, *args, **kwargs) -> typing.Any: + if self.session.model: + kwargs = {**kwargs, "model": self.session.model} + return fwd(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# Reporting the agent's activity as session/update notifications +# --------------------------------------------------------------------------- + + +class _PartialCall(typing.TypedDict): + """A tool call being assembled from streaming deltas.""" + + id: str + name: str + args: str + + +_TOOL_KINDS: dict[str, acp.schema.ToolKind] = { + # The harness's own tools for running model-authored Python. + StatefulReplSynthesizer.exec_code.__name__: "execute", + FinalBodySynthesizer._SubmitSolutionTool.__toolname__: "execute", + acp_read_text_file.__name__: "read", + acp_write_text_file.__name__: "edit", + acp_run_terminal_command.__name__: "execute", + # `acp_ask_user` is deliberately absent: ACP's `ToolKind` vocabulary has no entry + # for asking the user something, and `think` -- the agent reasoning -- is a + # different thing rather than a near fit. It gets no kind at all; see `_tool_kind`. +} + + +ASSUMED_CONTEXT_SIZE = 128_000 +"""What to assume a model's context window is when litellm has no entry for it. + +A guess, and the gauge it feeds is only as good as it. The alternative is no gauge at +all for any model litellm has not catalogued -- which is most new ones, and anything +behind a gateway -- and a roughly-right gauge is worth more to someone watching their +context fill than an empty space where one should be. +""" + + +@functools.cache +def _context_size(model: str) -> int: + """How many tokens of context `model` has. + + litellm's table is the only source this can consult, and it does not cover every + model. `ASSUMED_CONTEXT_SIZE` covers the rest, and the guess is said out loud + once, since a gauge drawn against the wrong denominator is worth knowing about. + + Cached because it is consulted after every completion and the answer never moves. + """ + try: + if size := int(litellm.get_model_info(model).get("max_input_tokens") or 0): + return size + except Exception: + pass + print( + f"note: litellm does not know the context size of {model!r}; the usage gauge " + f"assumes {ASSUMED_CONTEXT_SIZE:,} tokens", + file=sys.stderr, + ) + return ASSUMED_CONTEXT_SIZE + + +def _locations( + raw_input: collections.abc.Mapping[str, typing.Any], +) -> list[acp.schema.ToolCallLocation]: + """Which file a call is about, said in the field an editor reads for it. + + Editors attribute a turn's work to files -- "these three were edited" -- and offer + to jump to them. They will guess from the raw arguments if they must (Poolside's + looks for ``path``, ``file_path``, ``cwd`` and several more), but `locations` is + where the answer belongs, and it is the only one that can carry a line number. + + Derived from the arguments rather than declared per tool, so a tool added later + that takes a `path` is located without anyone remembering to do it. + """ + path = raw_input.get("path") + if not isinstance(path, str) or not path: + return [] + line = raw_input.get("line") + return [ + acp.schema.ToolCallLocation( + path=path, line=line if isinstance(line, int) else None + ) + ] + + +def _tool_kind(name: str) -> acp.schema.ToolKind | None: + """The ACP category an editor uses to pick an icon for a tool call, if it is known. + + Keyed by the name a tool is *advertised* under, since that is the name that comes + back in the model's reply. Every key above is read off the tool itself rather than + written out, because a key that drifts is invisible: the call still runs, the + editor just draws the wrong icon. + + Names are assigned per request (`_advertised_names`) and may be disambiguated, so + this is a lookup with a fallback rather than a table that has to be exhaustive. + + `None` rather than ``"other"`` for the fallback, though both are spelled the same + way in the protocol's own vocabulary. ``kind`` is optional in every message that + carries it, and omitting it is how ACP says nothing is claimed; ``"other"`` is a + positive claim about a call, and editors treat it as one -- Poolside gives it the + terminal icon, the same one it gives ``execute``, so an unclassified call is drawn + as a shell command. That is how `acp_ask_user`, which runs no commands at all, came + to look like one. + """ + return _TOOL_KINDS.get(name) + + +@dataclasses.dataclass +class ACPSessionReporter(ObjectInterpretation): + """Translate the agent's activity into `session/update` notifications. + + Forces `completion` onto the streaming path so the editor sees text as it is + produced rather than in one block at the end, and brackets every `call_tool` with + the status transitions an editor renders as a tool-call row. + """ + + session: ACPSession + + _open: dict[str, str] = dataclasses.field(default_factory=dict) + """Tool calls announced to the editor and not yet given a terminal status.""" + + _terminals: dict[str, list[str]] = dataclasses.field(default_factory=dict) + """Terminals a call has opened, which are how that call renders. See `_content`.""" + + _diffs: dict[str, list[ToolCallContent]] = dataclasses.field(default_factory=dict) + """Edits a call has made, likewise. See `_content`.""" + + running: str | None = None + """The id of the call currently executing, for a tool that wants to say so.""" + + finish_reason: str | None = None + """Why the *last* completion of this turn stopped, in the provider's vocabulary.""" + + tokens: collections.Counter = dataclasses.field(default_factory=collections.Counter) + + def begin_turn(self) -> None: + """Forget the last turn. Called by the server before the worker starts. + + The reporter outlives a turn -- it belongs to the session -- but everything it + counts is per-turn, and `usage` and `stop_reason` would otherwise report the + whole conversation's totals as this prompt's. + """ + self._open.clear() + self._terminals.clear() + self._diffs.clear() + self.running = None + self.finish_reason = None + self.tokens.clear() + + def _start(self, call_id: str, name: str, **kwargs) -> None: + """Announce a tool call, at most once per id.""" + if call_id in self._open: + return + self._open[call_id] = name + self.session.notify( + acp.start_tool_call(call_id, name, kind=_tool_kind(name), **kwargs) + ) + + def _finish( + self, call_id: str, status: acp.schema.ToolCallStatus, text: str + ) -> None: + """Give an announced call a terminal status, so the editor stops waiting.""" + self._open.pop(call_id, None) + self.session.notify( + acp.update_tool_call( + call_id, status=status, content=self._content(call_id, text) + ) + ) + self._terminals.pop(call_id, None) + self._diffs.pop(call_id, None) + + def _content(self, call_id: str, text: str) -> list[ToolCallContent]: + """How this call should render: as what it *did*, if that can be shown. + + A terminal or a diff is not one rendering among several. The editor streams a + terminal's output live and "continues to display it even after the terminal is + released", and it draws a diff as a before-and-after; the text the call also + produced is then the same information a second time -- and on the failing path + the error is in the terminal, which is where the user is already looking. + + `ToolCallUpdate.content` replaces the collection rather than appending to it, + so every update for such a call has to carry it again. Composing the content in + one place is what keeps that from being remembered at each call site. + """ + if terminals := self._terminals.get(call_id): + return [acp.tool_terminal_ref(terminal) for terminal in terminals] + if diffs := self._diffs.get(call_id): + return list(diffs) + return [acp.tool_content(acp.text_block(text))] + + def show_diff(self, path: str, new_text: str, old_text: str | None) -> None: + """Render the call now running as an edit to `path`. + + `old_text` may be `None` -- for a file being created, or an editor that would + not say what was there before. The editor then draws the new content alone, + which is less useful than a diff and still better than a line of prose saying + a write happened. + """ + if (call_id := self.running) is None: + return + self._diffs.setdefault(call_id, []).append( + acp.tool_diff_content(path, new_text, old_text) + ) + self.session.notify( + acp.update_tool_call(call_id, content=self._content(call_id, "")) + ) + + def show_terminal(self, terminal_id: str) -> None: + """Render the call now running as this live terminal. + + Called by `ACPToolRuntime` the moment the editor hands back a terminal id, so + the user watches the command run instead of watching a spinner. A no-op when + no call is running -- the tool is reachable outside a tool call, and a terminal + with no row to attach to is not an error. + """ + if (call_id := self.running) is None: + return + self._terminals.setdefault(call_id, []).append(terminal_id) + self.session.notify( + acp.update_tool_call(call_id, content=self._content(call_id, "")) + ) + + def abandon(self) -> None: + """Fail every call still open. Called once the turn is over, however it ended. + + A cancelled turn, or one whose skill raised, leaves calls the editor was told + had started and never hears about again -- rendered as a row that spins for + the rest of the session. ACP's `ToolCallStatus` has no ``cancelled``, so + ``failed`` is the only terminal status available to say so. + """ + for call_id in list(self._open): + self._finish(call_id, "failed", "the turn ended before this call finished") + + def stop_reason(self) -> acp.schema.StopReason: + """Why this turn ended, for a turn the model itself brought to a close. + + The two interesting answers are the ones an editor cannot infer: a reply cut + off at the token limit and a reply the provider refused both arrive as an + *answer*, and reporting either as `end_turn` tells the user their question was + answered when it was not. + """ + if self.finish_reason == "length": + return "max_tokens" + if self.finish_reason == "content_filter": + return "refusal" + return "end_turn" + + def usage(self) -> acp.schema.Usage | None: + """This turn's token counts, or `None` if no provider reported any.""" + if not self.tokens: + return None + return acp.schema.Usage( + input_tokens=self.tokens["prompt_tokens"], + output_tokens=self.tokens["completion_tokens"], + total_tokens=self.tokens["total_tokens"], + ) + + def _account(self, response: typing.Any) -> None: + """Record what one completion cost and why it stopped. + + Read off the value handed back rather than the request, so the streamed and + unstreamed paths are accounted identically: `litellm.stream_chunk_builder` + rebuilds both fields onto the response it assembles from the chunks. + """ + choices = getattr(response, "choices", None) or [] + if choices: + self.finish_reason = getattr(choices[0], "finish_reason", None) + if usage := getattr(response, "usage", None): + for field in ("prompt_tokens", "completion_tokens", "total_tokens"): + self.tokens[field] += getattr(usage, field, 0) or 0 + self._report_context(response) + + def _report_context(self, response: typing.Any) -> None: + """Say how full the model's context is, after each request. + + Not the same number as `usage`, which is what the whole turn *cost* and is + reported once at the end. This is how much room is left, reported as it fills, + and an editor draws it as a gauge -- the difference between "that turn was + expensive" and "you are nearly out of context". + + Everything here is skipped rather than guessed when it cannot be known. + `size` is required by the protocol, and a denominator nobody measured makes a + gauge that lies; litellm does not know every model, and the response is the + only place the model's real name can be read once a picker has changed it. + """ + usage = getattr(response, "usage", None) + used = getattr(usage, "prompt_tokens", 0) or 0 + model = getattr(response, "model", None) + if not used or not model: + return + size = _context_size(model) + cost = None + with contextlib.suppress(Exception): + if amount := litellm.completion_cost(completion_response=response): + cost = acp.schema.Cost(amount=float(amount), currency="USD") + self.session.notify( + acp.schema.UsageUpdate( + session_update="usage_update", used=used, size=int(size), cost=cost + ) + ) + + @implements(completion) + def completion(self, *args, **kwargs) -> typing.Any: + """Stream this request, reporting deltas as they arrive. + + Streaming is something this handler *adds* to a request that did not ask for + it, so it also owns the cost: a broken stream falls back to an ordinary + unstreamed request rather than failing a call that would have succeeded. The + retry is safe because a broken stream produced no result to duplicate. + + This is also the turn's meter -- every request the loop makes passes through + here exactly once -- so it is where what the last reply cost and why it + stopped are recorded, for `usage` and `stop_reason` to report. + """ + if self.session.cancel.is_set(): + raise SessionCancelled + + try: + response = self._streamed(*args, **kwargs) + except ( + litellm.exceptions.MidStreamFallbackError, + litellm.exceptions.APIConnectionError, + litellm.exceptions.Timeout, + ): + # Deliberately narrow: a refused request, a bad model name or a rejected + # response schema fails identically unstreamed, and re-issuing it would + # only pay for the same error twice. + response = fwd(*args, **kwargs) + self._account(response) + return response + + def _streamed(self, *args, **kwargs) -> typing.Any: + # `response_format` is None exactly when the skill returns `str` (see + # `call_assistant`). Any other answer is JSON shaped like the response + # format, and streaming it would show the editor a `{"value": ...}` wrapper + # being typed out; the decoded value is reported once, by the server. + is_prose = kwargs.get("response_format") is None + # `include_usage` is what makes the gauge report the provider's own numbers. + # Without it a stream carries no usage block at all, and the counts come from + # `stream_chunk_builder` tokenizing the request locally -- an estimate that + # cannot see cache reads or a provider's own accounting. Asking costs one extra + # chunk, whose `choices` are empty; the loop below appends before it skips + # those, so it still reaches the builder. A provider that does not understand + # the option has it dropped rather than refused, since the launcher sets + # `litellm.drop_params`. + stream = fwd( + *args, + **{ + "stream_options": {"include_usage": True}, + **kwargs, + "stream": True, + }, + ) + + # Asking for a stream does not guarantee getting one: an inner handler may + # answer from a cache or a fixture and hand back a settled response, ignoring + # the flag this handler added. Report that in one go rather than trying to + # iterate a response object. + if isinstance(stream, litellm.types.utils.ModelResponse): + return self._settled(stream, is_prose=is_prose) + + chunks: list[typing.Any] = [] + calls: dict[int, _PartialCall] = {} + for chunk in stream: + if self.session.cancel.is_set(): + raise SessionCancelled + chunks.append(chunk) + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta is None: + continue + + if delta.content and is_prose: + self.session.notify(acp.update_agent_message_text(delta.content)) + if reasoning := getattr(delta, "reasoning_content", None): + self.session.notify(acp.update_agent_thought_text(reasoning)) + + for fragment in delta.tool_calls or []: + slot = calls.setdefault( + fragment.index, {"id": "", "name": "", "args": ""} + ) + slot["id"] = getattr(fragment, "id", None) or slot["id"] + if function := getattr(fragment, "function", None): + slot["name"] = function.name or slot["name"] + slot["args"] += function.arguments or "" + if not slot["id"] or not slot["name"]: + continue + self._start(slot["id"], slot["name"], status="pending") + try: + raw_input = pydantic_core.from_json( + slot["args"], allow_partial="trailing-strings" + ) + except ValueError: + raw_input = None + self.session.notify( + acp.update_tool_call(slot["id"], raw_input=raw_input) + ) + + return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages")) + + def _settled( + self, response: litellm.types.utils.ModelResponse, *, is_prose: bool + ) -> typing.Any: + """Report a response that arrived whole, and hand it back unchanged.""" + choice = response.choices[0] + if not isinstance(choice, litellm.types.utils.Choices): + return response + if (content := choice.message.get("content")) and is_prose: + self.session.notify(acp.update_agent_message_text(content)) + if reasoning := choice.message.get("reasoning_content"): + self.session.notify(acp.update_agent_thought_text(reasoning)) + for raw in choice.message.get("tool_calls") or []: + self._start(str(raw.id), raw.function.name or "?", status="pending") + return response + + @implements(call_tool) + def call_tool(self, tool_call: DecodedToolCall) -> typing.Any: + """Bracket the call with the status transitions an editor renders. + + A failed call arrives here two ways, and both have to end as ``failed``. This + handler sits *above* `TenacityRetryer` -- `EffectfulACPAgent._answer` installs + the session's stack on top of the harness's -- and the retryer's whole job is + to turn a raising tool into that call's result, so on the ordinary path the + exception never reaches this ``except``: it comes back as a perfectly normal + return whose `result` is the error. Reporting on the exception alone would + show the user every failed call as completed, with the traceback rendered as + its output. + """ + if self.session.cancel.is_set(): + raise SessionCancelled + # A call announced while streaming carries only the name, since its + # arguments were still arriving; now that they are decoded, say what the + # call actually is. + self._start(tool_call.id, tool_call.name, status="pending") + raw_input = _raw_input(tool_call) + self.session.notify( + acp.update_tool_call( + tool_call.id, + status="in_progress", + title=_call_title(tool_call.name, raw_input), + raw_input=raw_input, + locations=_locations(raw_input), + ) + ) + # Named while it runs, so a tool that has something to show -- a terminal -- + # can find the row it belongs to. Restored rather than cleared, since a tool + # may itself call a Skill whose own tool calls nest inside this one. + outer, self.running = self.running, tool_call.id + try: + message, result, is_final = fwd(tool_call) + except ToolCallExecutionError as e: + self._finish(tool_call.id, "failed", str(e)) + raise + finally: + self.running = outer + + self._finish( + tool_call.id, + "failed" if isinstance(result, ToolCallExecutionError) else "completed", + _as_text(message), + ) + return (message, result, is_final) + + +def _as_text(message: typing.Any) -> str: + """A message's content as a string, however it was encoded. + + `~effectful.handlers.llm.harness.hooks.call_tool` encodes a result into content + blocks, so it may be a list rather than a string -- an image tool returns one. + + A *missing* content is the empty string, not the JSON below. An assistant turn + that only called tools has ``content: None``, and rendering that as ``"null"`` put + the literal word into the editor for every such turn a reloaded session replayed. + """ + content = message.get("content") if hasattr(message, "get") else None + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + _part_as_text(part) for part in content if isinstance(part, dict) + ) + return json.dumps(content, default=str) + + +def _part_as_text(part: collections.abc.Mapping[str, typing.Any]) -> str: + """One content block as text: its own if it has any, else a note that it exists. + + A tool result is not always text. `call_tool` encodes one that returns an image as + an ``image_url`` block, and reading only the ``text`` key across the blocks turned + that into the empty string -- so the editor rendered the call as having produced + nothing at all, which is the same thing it renders for a tool that printed + nothing. The model still receives the image either way; this is what stands in for + it on the screen. + + A placeholder rather than the block itself because the content here is bound for a + text block. Handing the editor a real `acp.image_block` is a change to + `ACPSessionReporter._content`, and worth making for a client that renders one -- + VS Code's does not, and would show nothing where this shows ``[image/png]``. + """ + if (kind := part.get("type")) == "text": + return part.get("text") or "" + if kind == "image_url": + url = part.get("image_url") + url = url.get("url", "") if isinstance(url, dict) else (url or "") + media = url[len("data:") :].split(";", 1)[0] if url.startswith("data:") else "" + return f"[{media or 'image'}]" + return f"[{kind or 'attachment'}]" + + +# --------------------------------------------------------------------------- +# Asking the editor's user before running a tool +# --------------------------------------------------------------------------- + + +def _raw_input(tool_call: DecodedToolCall) -> dict[str, typing.Any]: + """A call's arguments as the data an editor renders, back from their decoded form. + + Round-tripped through the encoding the model was given rather than read off + `bound_args`, so what the editor is shown is what the model actually said -- a + code object comes back as its source, an image as its reference. + """ + return json.loads( + pydantic.TypeAdapter(Encodable[DecodedToolCall]).dump_python( + tool_call, mode="json", context={} + )["function"]["arguments"] + ) + + +def _call_title(name: str, raw_input: collections.abc.Mapping[str, typing.Any]) -> str: + """A one-line description of a call: what it is, and what it was given. + + The bare tool name is not enough, and not only because it is terse. Clients treat + a title that is *only* an identifier as a placeholder and replace it with phrasing + of their own -- Poolside tests it against ``/^[a-z][a-z0-9_]*$/`` and, on a match, + ignores it -- so `exec_code` is discarded and rendered as "Run exec_code", with the + code nowhere on screen. A title with arguments in it survives that test. + + It also has to carry the arguments because nothing else reliably does. `rawInput` + is sent on every call, but a client that has tool *output* to show may prefer it: + Poolside renders a command, then content, then raw input, whichever comes first, + and every call of ours ends with content. So this is where the arguments are. + """ + return f"{name}({', '.join(f'{key}={_abbreviate(value)}' for key, value in raw_input.items())})" + + +def _abbreviate(value: typing.Any, limit: int = 60) -> str: + """One argument, short enough to sit in a title.""" + text = value if isinstance(value, str) else json.dumps(value, default=str) + text = " ".join(text.split()) + return text if len(text) <= limit else f"{text[:limit]}…" + + +@dataclasses.dataclass +class ACPPermissionGate(ObjectInterpretation): + """Ask the editor's user to approve each tool call before it runs. + + A handler rather than a tool, because unlike the editor capabilities above this is + not something the model calls -- it gates *every* tool call, including the agent's + own and the harness's ``exec_code``. + """ + + session: ACPSession + + _standing: dict[str, bool] = dataclasses.field(default_factory=dict) + + permission_options: typing.ClassVar[ + collections.abc.Sequence[acp.schema.PermissionOption] + ] = ( + acp.schema.PermissionOption( + option_id="allow_once", name="Allow", kind="allow_once" + ), + acp.schema.PermissionOption( + option_id="allow_always", name="Always allow", kind="allow_always" + ), + acp.schema.PermissionOption( + option_id="reject_once", name="Reject", kind="reject_once" + ), + acp.schema.PermissionOption( + option_id="reject_always", name="Always reject", kind="reject_always" + ), + ) + + @implements(call_tool) + def call_tool[T](self, tool_call: DecodedToolCall[T]) -> T: + """Run the call if it is approved; otherwise make the tool itself refuse. + + A `Tool` is an `Operation`, and `call_tool` invokes it, so handling it is all + a refusal takes: the call then fails the way any raising tool fails, and + `TenacityRetryer` reports it to the model as that call's result. Forwarding + rather than answering here is what keeps `HistoryBuilder` in the loop, so the + declined call is still answered and the conversation stays sendable. + """ + with handler({tool_call.tool: self._decide(tool_call)}): + return fwd(tool_call) + + def _decide[T]( + self, tool_call: DecodedToolCall[T] + ) -> collections.abc.Callable[..., T]: + """What should run for this call: the tool itself, or a refusal in its place. + + The session's mode is consulted before the user is, since a mode is the user + having answered these prompts in advance -- that is the whole of what picking + one means. `UNGATED_TOOLS` comes before even that, for the one tool whose + whole purpose is to ask the user something. + + Raises: + SessionCancelled: If the user dismissed the prompt instead of answering it. + """ + + def refused( + exc: Exception, *args: typing.Any, **kwargs: typing.Any + ) -> typing.NoReturn: + raise exc + + # Before the mode, because this is not a decision a mode makes: `UNGATED_TOOLS` + # is about a tool whose only effect is to ask the user, and Plan mode is as + # entitled to ask as Auto is. + if tool_call.name in UNGATED_TOOLS: + return lambda *a, **k: fwd() + if self.session.mode_id == AUTO: + return lambda *a, **k: fwd() + if self.session.mode_id == PLAN and tool_call.name in MUTATING_TOOLS: + return functools.partial( + refused, + PermissionError( + f"The call to `{tool_call.name}` did not run: this session is in " + f"Plan mode, which changes nothing in the user's editor. Say what " + f"you would do and why; the user can switch to Ask or Auto mode if " + f"they want it done." + ), + ) + + standing = self._standing.get(tool_call.name) + if standing is True: + return lambda *a, **k: fwd() + if standing is False: + return functools.partial( + refused, + PermissionError( + f"The call to `{tool_call.name}` did not run: the user declined it earlier. Do not retry it; either continue without it, or explain what you cannot do and why." + ), + ) + + raw_input = _raw_input(tool_call) + # No bound on this wait, deliberately: it is a dialog in front of a person. + # See `ACPSession.call`. + response = self.session.call( + self.session.client.request_permission( + self.session.session_id, + acp.schema.ToolCallUpdate( + tool_call_id=tool_call.id, + title=_call_title(tool_call.name, raw_input), + kind=_tool_kind(tool_call.name), + raw_input=raw_input, + ), + options=list(self.permission_options), + ) + ) + + outcome = response.outcome + # A `cancelled` outcome is the user dismissing the prompt, not rejecting the + # call; the turn is over either way. + if outcome.outcome != "selected": + raise SessionCancelled + + allowed = outcome.option_id.startswith("allow") + if outcome.option_id.endswith("always"): + self._standing[tool_call.name] = allowed + if allowed: + return lambda *a, **k: fwd() + return functools.partial( + refused, + PermissionError( + f"The call to `{tool_call.name}` did not run: the user declined it. Do not retry it; either continue without it, or explain what you cannot do and why." + ), + ) + + +# --------------------------------------------------------------------------- +# Remembering that a session existed +# --------------------------------------------------------------------------- + + +class SessionIndex: + """The sessions this agent knows of, and enough about each one to list it. + + ACP lets a client ask the agent what conversations it has (`session/list`), which + is how an editor fills a session picker that survives a restart. Answering needs + more than the agent histories `SQLitePersister` already keeps: `SessionInfo` + requires the `cwd` a session was opened on, and a useful listing wants a title and + a time. That is what this table holds. + + It lives *in the persistence database* rather than a file of its own, and exists + only when persistence does. Both follow from the same observation: a session + listed here whose history is not there would be an entry for a conversation that + cannot be reopened. So when no persistence handler is installed there is no index, + `session/list` is not advertised, and it is not answered -- which is deliberately + not the same as answering "no sessions". A client reconciles its own history + against this reply (VS Code's calls `reconcileFromAgent` with the ids it gets + back), so an empty answer tells it to forget every session it knew about. + """ + + SCHEMA: typing.ClassVar[str] = """ + CREATE TABLE IF NOT EXISTS acp_sessions ( + session_id TEXT PRIMARY KEY, + cwd TEXT NOT NULL, + additional_directories TEXT NOT NULL DEFAULT '[]', + title TEXT, + updated_at TEXT NOT NULL + ) + """ + + @classmethod + def open(cls) -> sqlite3.Connection | None: + """A connection to the index, or `None` if nothing is persisting anything. + + `SQLitePersister` hands back a fresh connection per call and says that is what + makes it safe from any thread, so this does not hold one. + """ + conn = SQLitePersister._checkpoint_connection() + if conn is not None: + with conn: + conn.execute(cls.SCHEMA) + return conn + + @classmethod + def available(cls) -> bool: + """Whether there is an index to answer from. Decides what is advertised.""" + return cls.open() is not None + + @classmethod + def record(cls, session: ACPSession) -> None: + """Note that this session exists, where it is rooted, and that it just moved. + + Called whenever a session is opened or answers a prompt, so `updated_at` + orders the listing by when each conversation was last used -- which is the + order a session picker wants. + """ + conn = cls.open() + if conn is None: + return + with conn: + conn.execute( + """ + INSERT INTO acp_sessions + (session_id, cwd, additional_directories, title, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(session_id) DO UPDATE SET + cwd = excluded.cwd, + additional_directories = excluded.additional_directories, + title = COALESCE(excluded.title, acp_sessions.title), + updated_at = excluded.updated_at + """, + ( + session.session_id, + session.cwd, + json.dumps(list(session.additional_directories)), + session.title or None, + datetime.datetime.now(datetime.UTC).isoformat(), + ), + ) + + @classmethod + def page( + cls, cwd: str | None, cursor: str | None, limit: int + ) -> tuple[list[acp.schema.SessionInfo], str | None]: + """One page of sessions, newest first, and the cursor for the next. + + Paged by key rather than by offset: the cursor names the last row handed out, + so a session that is written while the user is paging cannot push a row across + a page boundary and hide it. The cursor is opaque to the client, which is why + it can be this -- the two ordering columns, joined. + + Raises: + RequestError: If there is no index to read (see the class docstring). + """ + conn = cls.open() + if conn is None: + raise acp.RequestError.invalid_request( + { + "reason": ( + "this agent keeps no session index; it was started without a " + "persistence handler, so its sessions end with the process" + ) + } + ) + where, params = ["1 = 1"], [] + if cwd is not None: + where.append("cwd = ?") + params.append(cwd) + if cursor is not None: + updated_at, _, session_id = cursor.partition("\x1f") + where.append("(updated_at, session_id) < (?, ?)") + params += [updated_at, session_id] + rows = conn.execute( + f"SELECT session_id, cwd, additional_directories, title, updated_at " # noqa: S608 + f"FROM acp_sessions WHERE {' AND '.join(where)} " + f"ORDER BY updated_at DESC, session_id DESC LIMIT ?", + (*params, limit + 1), + ).fetchall() + more = len(rows) > limit + rows = rows[:limit] + sessions = [ + acp.schema.SessionInfo( + session_id=session_id, + cwd=cwd_, + additional_directories=json.loads(directories), + title=title, + updated_at=updated_at, + ) + for session_id, cwd_, directories, title, updated_at in rows + ] + next_cursor = f"{rows[-1][4]}\x1f{rows[-1][0]}" if more and rows else None + return sessions, next_cursor + + +def _title_from(text: str) -> str: + """A session title taken from the first thing the user said in it. + + The alternative is asking the model for one, which costs a request and a wait + before the answer the user is actually waiting for. Their own opening line is + usually what they would have called it anyway. + """ + line = " ".join(text.split()) + return line if len(line) <= 60 else line[:59].rstrip() + "…" + + +# --------------------------------------------------------------------------- +# The server +# --------------------------------------------------------------------------- + + +def _replay( + history: collections.abc.Iterable[collections.abc.Mapping[str, typing.Any]], +) -> collections.abc.Iterator[typing.Any]: + """The `session/update` notifications that reproduce a stored conversation. + + ACP requires `session/load` to replay "the entire conversation", and a coding + agent's conversation is mostly not prose: an assistant turn that read three files + carries no text at all, only tool calls, and dropping those would replay a + conversation in which the agent sat silent and then knew things. So each stored + tool call comes back as a completed tool-call row, and each stored tool *result* + fills in that row's output. + + A generator rather than a method, so it can be read against a history without a + session, an editor, or an event loop. + """ + for message in history: + role, text = message.get("role"), _as_text(message) + if role == "user": + if text: + yield acp.update_user_message_text(text) + elif role == "assistant": + if text: + yield acp.update_agent_message_text(text) + for raw in message.get("tool_calls") or []: + function = raw.get("function") or {} + name = function.get("name") or "?" + arguments = function.get("arguments") + try: + raw_input = ( + json.loads(arguments) + if isinstance(arguments, str) + else arguments + ) + except ValueError: + raw_input = None + yield acp.start_tool_call( + str(raw.get("id")), + name, + kind=_tool_kind(name), + # Completed, because a stored call is one that already ran: the + # turn it belonged to is over, whatever became of the call. + status="completed", + raw_input=raw_input, + ) + elif role == "tool" and (call_id := message.get("tool_call_id")) is not None: + yield acp.update_tool_call( + str(call_id), + content=[acp.tool_content(acp.text_block(text))], + ) + + +@dataclasses.dataclass(frozen=True) +class SlashCommand: + """One ``/name`` command: what the editor is told about it, and what runs it. + + The pairing is the point. A command exists in two conversations -- it is + *advertised* (`available_commands_update`, so the editor can offer it while the + user types) and it is *dispatched* (`EffectfulACPAgent._command`, when a prompt + arrives spelling it) -- and holding both halves in one value is what makes + those agree by construction. Kept as separate lists, adding a command to one + and forgetting the other would produce an editor offering something answered + with "Unknown command", and nothing anywhere would fail. + """ + + spec: acp.schema.AvailableCommand + """The advertisement: name, description, and the input hint, if any.""" + + run: collections.abc.Callable[["EffectfulACPAgent", ACPSession, str], str] + """The behaviour: handed the server, the session and the argument text.""" + + +def _run_clear(server: "EffectfulACPAgent", session: ACPSession, argument: str) -> str: + session.agent.__history__.clear() + return "Cleared. I have forgotten the conversation up to here." + + +def _run_status(server: "EffectfulACPAgent", session: ACPSession, argument: str) -> str: + roots = "\n".join(f"- `{root}`" for root in session.roots) or "- (none)" + mode = next( + (m.name for m in SESSION_MODES if m.id == session.mode_id), + session.mode_id, + ) + model = session.model or "as configured at launch" + return ( + f"**Mode** {mode}\n\n**Model** {model}\n\n" + f"**Directories**\n{roots}\n\n" + f"**Messages so far** {len(session.agent.__history__)}" + ) + + +def _run_mode(server: "EffectfulACPAgent", session: ACPSession, mode_id: str) -> str: + """``/mode`` with no argument reports; with one, switches. + + The editor's own picker sends `session/set_config_option`, and this sends + nothing -- it is already inside the agent. What it must do instead is *say* the + mode changed, on both channels a client might be listening to: + `current_mode_update` for one that reads `modes`, and the config options for + one that reads those. + """ + offered = {mode.id: mode for mode in SESSION_MODES} + if not mode_id: + return "\n".join( + [f"**Mode** {offered[session.mode_id].name}", ""] + + [f"- `/mode {mode.id}` — {mode.description}" for mode in SESSION_MODES] + ) + if mode_id not in offered: + return ( + f"No such mode `{mode_id}`. Try " + f"{', '.join(f'`{mode_id}`' for mode_id in offered)}." + ) + session.mode_id = mode_id + session.notify( + acp.schema.CurrentModeUpdate( + session_update="current_mode_update", current_mode_id=mode_id + ) + ) + session.notify( + acp.schema.ConfigOptionUpdate( + session_update="config_option_update", + config_options=server._config_options(session), + ) + ) + return f"Mode is now **{offered[mode_id].name}**. {offered[mode_id].description}" + + +_SLASH_COMMANDS: dict[str, SlashCommand] = { + command.spec.name: command + for command in ( + SlashCommand( + spec=acp.schema.AvailableCommand( + name="clear", + description="Forget the conversation so far, keeping this session open.", + ), + run=_run_clear, + ), + SlashCommand( + spec=acp.schema.AvailableCommand( + name="status", + description="Show the mode, model and directories this session is using.", + ), + run=_run_status, + ), + SlashCommand( + spec=acp.schema.AvailableCommand( + name="mode", + description="Switch how much this agent may do without asking.", + # A command may take an argument, and the hint is what the editor + # shows after the name while the user is typing it. Derived from + # `SESSION_MODES`, as everything mode-shaped is. + input=acp.schema.AvailableCommandInput( + acp.schema.UnstructuredCommandInput( + hint=" | ".join(mode.id for mode in SESSION_MODES) + ) + ), + ), + run=_run_mode, + ), + ) +} +"""Every command, dispatch and advertisement together. See `SlashCommand`.""" + +SLASH_COMMANDS: tuple[acp.schema.AvailableCommand, ...] = tuple( + command.spec for command in _SLASH_COMMANDS.values() +) +"""The advertised half of `_SLASH_COMMANDS`, in the shape the notification takes.""" + + +class Attachment(pydantic.BaseModel): + """A file or resource attached by reference URI""" + + uri: str + + @pydantic.field_validator("uri") + @classmethod + def _as_path(cls, uri: str) -> str: + """ + A ``file:`` URI becomes the plain path the read tool takes. + Anything else is passed through unchanged for the agent to interpret. + """ + parsed = urllib.parse.urlsplit(uri) + if parsed.scheme == "file": + return urllib.request.url2pathname(parsed.path) + return uri + + +def _prompt_parts( + prompt: list[ContentBlock], +) -> tuple[str, list[Attachment], list[Image.Image]]: + """Split a prompt into the arguments the agent's ``prompt`` skill takes. + + Two block kinds are baseline -- every agent must handle them, with no capability + to negotiate -- and one is claimed (``image``): + + * ``text``, the user's own words, joined into the prose argument. + * ``resource_link``, a *reference* to a file rather than its contents. It + becomes an `Attachment` -- the path, so the model can decide to open it with + `acp_read_text_file`, which reads through the editor and therefore sees + unsaved changes. Inlining a file here would freeze a stale copy into the + conversation and pay its tokens on every request after, whether or not the + answer ever needed it. + * ``image``, decoded into the `PIL.Image.Image` the skill accepts -- when it + carries its data. One that is only a URI is refused: nothing here fetches. + + Everything else is refused rather than dropped -- an audio clip, and notably + ``resource``, a file's contents inlined by the editor. This agent deliberately + does not claim `embedded_context`, and a conforming client then sends links + instead (ACP: capabilities not claimed are unsupported, and clients MUST + restrict prompt content accordingly), which is the whole point: the flat fee + for attaching a large file becomes one line, and reading it back is bounded + and on demand. A non-conforming block is refused so the client is told, not + quietly answered as though the attachment were never there. Silently + discarding an attachment is the failure mode that looks like success. + + Raises: + RequestError: If the prompt is empty, or carries a block this cannot read. + """ + + def unreadable(what: str) -> acp.RequestError: + return acp.RequestError.invalid_params( + {"reason": f"this agent cannot read {what}"} + ) + + texts: list[str] = [] + attachments: list[Attachment] = [] + images: list[Image.Image] = [] + for block in prompt: + if block.type == "text": + texts.append(block.text) + elif block.type == "resource_link": + attachments.append(Attachment(uri=block.uri)) + elif block.type == "image": + if not block.data: + raise unreadable(f"an image that is only a reference ({block.uri})") + images.append(Image.open(io.BytesIO(base64.b64decode(block.data)))) + else: + raise unreadable( + f"a {block.type!r} block; it advertises no prompt capability for one" + ) + text = "".join(texts).strip() + if not text and not attachments and not images: + raise acp.RequestError.invalid_params({"reason": "the prompt is empty"}) + return text, attachments, images + + +class EffectfulACPAgent[A: Agent](acp.Agent): + """An ACP server backed by one `Agent` instance per session. + + Parameterised by how to *make* an agent rather than by an agent, so this module + never has to know about any particular one. `make_agent` is handed the session id + and returns an agent bearing it as its ``__agent_id__``; an agent class with an + ``__agent_id__`` field is already such a callable, which is the usual way to pass + one (see ``assistant.py``). + + The agent must have a ``prompt`` skill -- named for the protocol method it + answers, ``session/prompt`` -- of the form:: + + prompt(user_input: str, + attachments: Sequence[Attachment] = (), + images: Sequence[Image.Image] = ()) -> ... + + The contract is assumed, not discovered: `_answer` calls it directly, and the + capabilities advertised below claim exactly what it accepts. Introspecting each + agent for what it happens to take would make the advertisement -- sent once at + ``initialize`` -- depend on an agent that does not exist yet. + + `models` fills the editor's model picker, and defaults to reading `OFFER_MODELS_ENV` + rather than to nothing. Which side of this module that default lives on is the + whole question: an editor configures an agent with a command and an environment, so + it is the *server* that knows to look there, not the script it is serving. Leaving + it to the caller would put a few lines of environment parsing in every script that + wanted a picker, and each of them would be a chance to spell it differently. + """ + + make_agent: collections.abc.Callable[[str], A] + models: tuple[str, ...] + page_size: int + + client: acp.interfaces.Client + client_capabilities: acp.schema.ClientCapabilities + + def __init__( + self, + make_agent: collections.abc.Callable[[str], A], + *, + models: collections.abc.Sequence[str] | None = None, + page_size: int = 50, + ): + self.make_agent = make_agent + # `None` rather than `()` as the default, because "the caller said nothing" and + # "the caller said no models" are different answers and only the first should + # consult the environment. A caller passing `()` has turned the picker off. + self.models = _offered_models() if models is None else tuple(models) + self.page_size = page_size + self.sessions: dict[str, ACPSession[A]] = {} + + @property + def agent_capabilities(self) -> acp.schema.AgentCapabilities: + """The capabilities this agent advertises to the editor. + + The editor uses them to decide what to offer the user, and the model uses + them to decide what to ask the editor to do. Everything claimed here is + something implemented below, and the reverse also has to hold: a client "MUST + verify that the Agent supports this capability" before using one, so a method + this class defines but does not advertise is a method no conforming client + will ever call. `close_session` was exactly that until this said so, which + left every session and its writer task alive for the life of the process. + + `prompt_capabilities` claims exactly what the ``prompt`` skill contract + accepts -- see `initialize` for the argument. `mcp_*` is left claiming + nothing, which is the honest answer for an agent that connects to no MCP + servers. + """ + return acp.schema.AgentCapabilities( + load_session=True, + prompt_capabilities=acp.schema.PromptCapabilities(image=True), + session_capabilities=acp.schema.SessionCapabilities( + close=acp.schema.SessionCloseCapabilities(), + resume=acp.schema.SessionResumeCapabilities(), + fork=acp.schema.SessionForkCapabilities(), + # Conditional, because this one is a claim about *state*: without a + # persistence handler there are no sessions to list, and saying + # otherwise invites a client to ask a question with no good answer. + list=acp.schema.SessionListCapabilities() + if SessionIndex.available() + else None, + ), + ) + + def _modes(self, session: ACPSession[A]) -> acp.schema.SessionModeState: + """The modes on offer and the one in force, for a session response.""" + return acp.schema.SessionModeState( + current_mode_id=session.mode_id, available_modes=list(SESSION_MODES) + ) + + def _config_options(self, session: ACPSession[A]) -> list[ConfigOption]: + """Every control this session puts in the editor's UI. + + The mode is here *as well as* in `modes`, which looks like saying it twice and + is not. A client that understands config options "MUST use them exclusively + and ignore the legacy modes field" -- so the moment this list is non-empty, a + client that reads it hides its mode picker and looks for an option whose + category is ``mode`` instead. Offering only the model would therefore take the + mode picker away from exactly the clients that render pickers best. `modes` + stays in the response for clients that do not read this list at all. + + The model option appears only when this server was given models to choose + between: an option listing one choice is a control that does nothing, which is + worse in a user interface than no control. + """ + options: list[ConfigOption] = [ + acp.schema.SessionConfigOptionSelect( + type="select", + id=MODE_OPTION_ID, + name="Mode", + description="How much this agent may do without asking.", + category="mode", + current_value=session.mode_id, + options=[ + acp.schema.SessionConfigSelectOption( + value=mode.id, name=mode.name, description=mode.description + ) + for mode in SESSION_MODES + ], + ) + ] + if self.models: + options.append( + acp.schema.SessionConfigOptionSelect( + type="select", + id=MODEL_OPTION_ID, + name="Model", + description="Which model answers in this session.", + category="model", + current_value=session.model, + options=[ + acp.schema.SessionConfigSelectOption( + value=INHERIT_MODEL, + name="Default", + description="Whatever this agent process was started with.", + ), + *( + acp.schema.SessionConfigSelectOption( + value=model, name=model + ) + for model in self.models + ), + ], + ) + ) + return options + + def _announce_commands(self, session: ACPSession[A]) -> None: + """Tell the editor which ``/name`` commands to offer for this session. + + Sent as a notification once the session exists, which is what the spec + describes. It races the response to the request that created the session -- + both go out on one pipe from two tasks -- and a client that has not yet learned + the id may drop it. Nothing is lost that matters: the commands are a + convenience, and reopening the session announces them again. + """ + session.notify( + acp.schema.AvailableCommandsUpdate( + session_update="available_commands_update", + available_commands=list(SLASH_COMMANDS), + ) + ) + + @property + def agent_info(self) -> acp.schema.Implementation: + """The agent's name, title and version, for the editor to display.""" + return acp.schema.Implementation( + name="effectful", title="effectful.handlers.llm", version="0.4.0" + ) + + def _open_session( + self, + session_id: str, + cwd: str, + additional_directories: list[str] | None, + ) -> ACPSession[A]: + """Open this session, or re-point an already open one at these directories. + + Opening on demand is what makes `load_session` work at all: after a restart + the editor knows a session id that this process has never seen, and the agent + constructed under it reads its own history back from the checkpoint. Only + ``session/new`` and ``session/load`` may do it, though -- see `_session`. + + Call it from the event loop thread, since `ACPSession` starts a task there. + """ + roots = tuple(additional_directories or ()) + if session_id not in self.sessions: + self.sessions[session_id] = ACPSession( + agent=self.make_agent(session_id), + client=self.client, + client_capabilities=self.client_capabilities, + cwd=cwd, + additional_directories=roots, + ) + else: + # An editor may reopen a session it still has open, and may do so from a + # different window onto a different directory. The conversation is the + # same one; where it is rooted is whatever it was just told. + session = self.sessions[session_id] + session.cwd, session.additional_directories = cwd, roots + return self.sessions[session_id] + + def _session(self, session_id: str) -> ACPSession[A]: + """This session, which must already be open. + + Every method other than ``session/new`` and ``session/load`` names a session + the editor believes is open, so an id that is not is a mistake and is answered + as one. Opening one here instead would turn a typo -- or a prompt sent against + a session that was never loaded -- into a silently fresh conversation, which + looks to the user like an agent that forgot everything. + + Raises: + RequestError: If no session is open under `session_id`. + """ + session = self.sessions.get(session_id) + if session is None: + raise acp.RequestError.resource_not_found(session_id) + return session + + def _decline_mcp(self, mcp_servers: list[typing.Any] | None) -> None: + """Note, without refusing, that this agent will not use the editor's MCP servers. + + Agents "SHOULD connect to all MCP servers specified by the Client", and stdio + transport is baseline -- there is no capability with which to say "none at + all", so a client with servers configured will send them on every + ``session/new`` and is behaving correctly in doing so. Failing the request + over that would make this agent unusable in any editor that has an MCP server + set up, to no one's benefit; ignoring them silently would hide it. stderr is + free (`serve` gives the protocol its own descriptor), so it goes there. + """ + if mcp_servers: + print( + f"note: ignoring {len(mcp_servers)} MCP server(s) offered by the " + f"editor; this agent has no MCP client", + file=sys.stderr, + ) + + def on_connect(self, conn: acp.interfaces.Client) -> None: + self.client = conn + + async def initialize( + self, + protocol_version: int, + client_capabilities: acp.schema.ClientCapabilities | None = None, + client_info: acp.schema.Implementation | None = None, + **kwargs: typing.Any, + ) -> acp.schema.InitializeResponse: + """Negotiate: say what this agent can do, and remember what the client can. + + `prompt_capabilities` claims what `_prompt_parts` reads and nothing more -- + the contract: a client "MUST adapt its interface according to + `PromptCapabilities`", and treats anything not claimed as unsupported. Both + directions of that rule are used deliberately here: + + * ``image`` is claimed, because the ``prompt`` skill contract takes decoded + images -- a promise that an attached screenshot will be *looked at*. + * ``embedded_context`` is not, and its absence is load-bearing: it is what + makes a conforming client attach a file as a ``resource_link`` -- a + reference costing a line -- instead of inlining its whole contents into + a prompt this agent would then be carrying in the conversation, and + paying for, on every request after. The model reads an attachment + through `acp_read_text_file` if and when the request needs it, bounded + and fresh from the editor's buffer. (Poolside's client, for one, + auto-attaches the active file to every prompt with its full text when + this is claimed, and degrades to links itself when it is not.) + + The client's own capabilities are consulted by `ACPToolRuntime`: the editor + tools are offered to the model either way, and one the client cannot service + reports itself as a failed call rather than being withheld. + """ + self.client_capabilities = ( + client_capabilities or acp.schema.ClientCapabilities() + ) + return acp.schema.InitializeResponse( + protocol_version=acp.PROTOCOL_VERSION, + agent_capabilities=self.agent_capabilities, + agent_info=self.agent_info, + ) + + async def new_session( + self, + cwd: str, + additional_directories: list[str] | None = None, + mcp_servers: list[typing.Any] | None = None, + **kwargs: typing.Any, + ) -> acp.schema.NewSessionResponse: + """Open a session, and with it a fresh agent. + + The session id becomes the agent's ``__agent_id__``, which is what makes the + conversation persistent: with a persistence handler installed, the agent's + history and declared fields are checkpointed under that id after every call + and restored by `load_session` below. + + `cwd` is the directory the user opened, and the session keeps it: it is what + the model is told it is working on, and where a terminal command runs. + """ + self._decline_mcp(mcp_servers) + session = self._open_session(str(uuid.uuid4()), cwd, additional_directories) + SessionIndex.record(session) + self._announce_commands(session) + return acp.schema.NewSessionResponse( + session_id=session.session_id, + modes=self._modes(session), + config_options=self._config_options(session), + ) + + async def load_session( + self, + cwd: str, + session_id: str, + mcp_servers: list[typing.Any] | None = None, + additional_directories: list[str] | None = None, + **kwargs: typing.Any, + ) -> acp.schema.LoadSessionResponse: + """Reopen an earlier session and replay it to the editor. + + Constructing the agent under the same id is the whole of the restore: + `Agent.__history__` reads the checkpoint lazily on first use. The replay is + required -- the agent MUST stream the *entire* conversation back, and MUST + wait until it has, because the client may be a different process with no + other record of it. + """ + self._decline_mcp(mcp_servers) + session = self._open_session(session_id, cwd, additional_directories) + SessionIndex.record(session) + for update in _replay(session.agent.__history__): + session.notify(update) + self._announce_commands(session) + await session.flush() + return acp.schema.LoadSessionResponse( + modes=self._modes(session), + config_options=self._config_options(session), + ) + + async def resume_session( + self, + session_id: str, + cwd: str, + additional_directories: list[str] | None = None, + mcp_servers: list[typing.Any] | None = None, + **kwargs: typing.Any, + ) -> acp.schema.ResumeSessionResponse: + """Reopen a session the client already has the transcript of. + + The same restore as `load_session` without the replay: a client resumes rather + than loads exactly when it kept its own copy of the conversation and wants the + agent to pick it up, not to be told it again. That is the whole difference, and + it is why both are advertised -- a client picks one. + """ + self._decline_mcp(mcp_servers) + session = self._open_session(session_id, cwd, additional_directories) + SessionIndex.record(session) + self._announce_commands(session) + return acp.schema.ResumeSessionResponse( + modes=self._modes(session), + config_options=self._config_options(session), + ) + + async def fork_session( + self, + session_id: str, + cwd: str, + additional_directories: list[str] | None = None, + mcp_servers: list[typing.Any] | None = None, + **kwargs: typing.Any, + ) -> acp.schema.ForkSessionResponse: + """Branch a conversation: a new session that starts as a copy of this one. + + For trying a second approach without losing the first. The copy is a copy -- + a new id, a new agent, its own history that happens to begin the same way -- + so a turn in either leaves the other alone. The user's settings come with it, + since a fork is a continuation and re-picking the mode and model would be a + chore rather than a choice. + + Reading the source through `_open_session` rather than requiring it open lets + a session be forked from a list after a restart, when its history lives only + in the checkpoint. The fork's own history is checkpointed when it first + answers, as any session's is -- so a fork abandoned before its first turn is + an empty conversation, not a copy. + """ + self._decline_mcp(mcp_servers) + source = self._open_session(session_id, cwd, additional_directories) + fork = self._open_session(str(uuid.uuid4()), cwd, additional_directories) + fork.agent.__history__.extend(source.agent.__history__) + fork.mode_id, fork.model = source.mode_id, source.model + fork.title = f"{source.title} (fork)" if source.title else "" + SessionIndex.record(fork) + self._announce_commands(fork) + return acp.schema.ForkSessionResponse( + session_id=fork.session_id, + modes=self._modes(fork), + config_options=self._config_options(fork), + ) + + async def list_sessions( + self, + cwd: str | None = None, + cursor: str | None = None, + **kwargs: typing.Any, + ) -> acp.schema.ListSessionsResponse: + """The conversations this agent knows of, newest first. + + This is what lets an editor offer sessions from before it was last closed. + `cwd` narrows the answer to one project, which is what a client asks for when + its window is that project. + + Raises: + RequestError: If this agent keeps no session index (see `SessionIndex`). + """ + sessions, next_cursor = SessionIndex.page(cwd, cursor, self.page_size) + return acp.schema.ListSessionsResponse( + sessions=sessions, next_cursor=next_cursor + ) + + async def set_session_mode( + self, session_id: str, mode_id: str, **kwargs: typing.Any + ) -> acp.schema.SetSessionModeResponse: + """Switch this session's mode, for a client that does not read config options. + + Raises: + RequestError: If the session is not open, or the mode is not one offered. + """ + self._set_mode(self._session(session_id), mode_id) + return acp.schema.SetSessionModeResponse() + + async def set_config_option( + self, config_id: str, session_id: str, value: str | bool, **kwargs: typing.Any + ) -> acp.schema.SetSessionConfigOptionResponse: + """Set one of this session's config options, and answer with all of them. + + The response is the whole list rather than an acknowledgement, because a + client redraws its controls from it -- which is also why setting an unknown + option is refused rather than ignored: a control that silently does nothing is + worse than one that reports it cannot. + + Raises: + RequestError: If the session, the option, or the value is not known. + """ + session = self._session(session_id) + if not isinstance(value, str): + raise acp.RequestError.invalid_params( + {"reason": f"{config_id!r} takes a string, not {value!r}"} + ) + if config_id == MODE_OPTION_ID: + self._set_mode(session, value) + elif config_id == MODEL_OPTION_ID and self.models: + if value != INHERIT_MODEL and value not in self.models: + raise acp.RequestError.invalid_params( + {"reason": f"{value!r} is not one of the models on offer"} + ) + session.model = value + else: + raise acp.RequestError.invalid_params( + {"reason": f"no such config option: {config_id!r}"} + ) + return acp.schema.SetSessionConfigOptionResponse( + config_options=self._config_options(session) + ) + + def _set_mode(self, session: ACPSession[A], mode_id: str) -> None: + """Switch a session's mode, however the editor asked for it. + + Both ways in end here: `session/set_mode`, and `session/set_config_option` on + the ``mode`` option, which is what a client that reads config options sends + instead. + + Raises: + RequestError: If `mode_id` is not one of `SESSION_MODES`. + """ + if mode_id not in {mode.id for mode in SESSION_MODES}: + raise acp.RequestError.invalid_params( + {"reason": f"no such mode: {mode_id!r}"} + ) + session.mode_id = mode_id + + async def prompt( + self, + session_id: str, + prompt: list[ContentBlock], + **kwargs: typing.Any, + ) -> acp.schema.PromptResponse: + """Answer one prompt, reporting progress as it goes. + + A prompt is a *list* of content blocks, not a string: the user's typed text, + plus whatever their editor attached -- referenced files, screenshots. + `_prompt_parts` splits it into the arguments the agent's ``prompt`` skill + takes, and rejects what this agent cannot read rather than dropping it + (see there). + + The skill call is synchronous and can run for minutes, so it goes to a worker + thread. `asyncio.to_thread` copies this task's context, and the handler stack + lives in a `ContextVar`, so the worker inherits the ambient stack -- the one + the module launcher installed -- and adds this session's own handlers to it. + + The lock is what keeps one session to one turn. Nothing upstream provides it: + `acp.connection` dispatches each request as its own task and does not await + it, so a client that sends a second prompt before the first has answered would + otherwise put two worker threads on one agent's history. + + The stop reason is the turn's summary, and the interesting values all come + from somewhere other than a normal return: `cancelled` is raised out of the + loop, and `max_tokens` and `refusal` are read by `ACPSessionReporter` off the + last reply. Only a turn with nothing else to say is `end_turn`. + """ + session = self._session(session_id) + async with session.lock: + session.cancel.clear() + session.reporter.begin_turn() + try: + text, attachments, images = _prompt_parts(prompt) + self._retitle(session, text) + if (answered := self._command(session, text)) is not None: + session.notify(acp.update_agent_message_text(answered)) + return acp.schema.PromptResponse(stop_reason="end_turn") + answer = await asyncio.to_thread( + self._answer, session, text, attachments, images + ) + # A `str` skill's answer was already streamed to the editor token by + # token; anything else was decoded from JSON that would have been noise + # to stream, so it is reported here, once, in its decoded form. + if not isinstance(answer, str): + session.notify( + acp.update_agent_message_text(json.dumps(answer, default=str)) + ) + stop_reason = session.reporter.stop_reason() + except SessionCancelled: + stop_reason = "cancelled" + finally: + # Whatever happened -- an answer, a cancellation, a skill that raised + # past this and out to the connection -- the editor is left with no + # tool call still spinning, and hears everything before it hears the + # result. + session.reporter.abandon() + await session.flush() + return acp.schema.PromptResponse( + stop_reason=stop_reason, usage=session.reporter.usage() + ) + + def _retitle(self, session: ACPSession[A], text: str) -> None: + """Name the session after its first prompt, and tell the editor the name. + + Once, on the first thing the user says: a title that followed the latest + message would rename the conversation out from under whoever is reading the + list. A slash command does not name a session either -- ``/status`` is not + what the conversation is about. + + Sessions are addressed by an opaque id, so without this a session list shows + the user a column of UUIDs. + """ + if session.title or text.startswith("/"): + SessionIndex.record(session) + return + session.title = _title_from(text) + SessionIndex.record(session) + session.notify( + acp.schema.SessionInfoUpdate( + session_update="session_info_update", + title=session.title, + updated_at=datetime.datetime.now(datetime.UTC).isoformat(), + ) + ) + + def _command(self, session: ACPSession[A], text: str) -> str | None: + """Answer `text` here if it is a slash command, or `None` to send it onward. + + A command is an ordinary prompt whose text begins with the name -- ACP has no + separate method for one -- so recognising the prefix and looking the name up + in `_SLASH_COMMANDS` is the whole mechanism; the same table is what + `_announce_commands` advertises, so a command offered is a command answered. + All of them run without a model: they are about the session rather than about + anything the model would know, and paying for a round trip to be told the + working directory would be an odd way to spend the user's money. + + Both go into the reply as prose rather than into the agent's history, so the + model never sees the exchange. `/clear` in particular must not: a message + saying the conversation was forgotten is the one thing that should not survive + forgetting it. + """ + if not text.startswith("/"): + return None + name, _, argument = text[1:].partition(" ") + name, argument = name.strip(), argument.strip() + command = _SLASH_COMMANDS.get(name) + if command is None: + return f"Unknown command `/{name}`. Try {', '.join(f'`/{c.name}`' for c in SLASH_COMMANDS)}." + return command.run(self, session, argument) + + def _answer( + self, + session: ACPSession[A], + text: str, + attachments: list[Attachment], + images: list[Image.Image], + ) -> typing.Any: + """Call the agent's ``prompt`` skill under this session's handlers. + + Runs in a worker thread. The call is direct rather than introspected: the + skill contract is this server's to define (see the class docstring), and a + nonconforming agent should fail loudly at its first prompt, the way any + wrong argument list does. + + Installing on top of the ambient stack, rather than assembling one, is what + lets the launcher decide the model, the retry budget and the persistence: the + session contributes only its three translations to the editor. + """ + with handler(session.intp): + # `Agent` the bound says nothing about a `prompt` skill; the contract + # is this server's own (class docstring), so the checker is waved off + # here rather than widened everywhere the type parameter travels. + return session.agent.prompt( # type: ignore + text, attachments=attachments, images=images + ) + + async def cancel(self, session_id: str, **kwargs: typing.Any) -> None: + """Ask the worker to stop at its next cancellation point. + + A notification, so it must not block: setting the flag is the whole of it, and + the turn reads it before the next completion, before the next tool call, + between stream chunks, and while waiting on the editor. + + A notification also has nowhere to report an error, so an id with no session + behind it is dropped rather than raised on -- and, unlike every other method + here, must not open one, since cancelling a session that does not exist would + otherwise create it. + """ + if session := self.sessions.get(session_id): + session.cancel.set() + + async def close_session( + self, session_id: str, **kwargs: typing.Any + ) -> acp.schema.CloseSessionResponse: + """Stop this session's turn and its writer, and forget it. + + Dropping the entry matters: the writer task does not survive being cancelled, + so a session left in the table after this would accept notifications that + nothing delivers, and the first turn to wait for its queue to empty would wait + forever. Forgetting it means a later `load_session` under the same id builds a + working one instead. + """ + session = self.sessions.pop(session_id, None) + if session is None: + return acp.schema.CloseSessionResponse() + session.cancel.set() + if session.writer is not None: + session.writer.cancel() + with contextlib.suppress(asyncio.CancelledError): + await session.writer + return acp.schema.CloseSessionResponse() + + async def serve(self) -> None: + """Serve one agent over stdio until the editor disconnects. + + stdout is the protocol, and the harness runs model-authored Python that may print + to it. So fd 1 is pointed at stderr for the process's lifetime, after handing a + duplicate of the real one to the transport -- which captures the file descriptor + when the streams are built and is unaffected by the later rebinding. + + No ``receive_timeout``, deliberately. That parameter bounds how long the + transport will wait for the *next message from the editor*, and tears the + connection down when it expires -- so any value at all is a rule that the user + may not think for longer than it before their agent disappears mid-conversation, + which is what a minute of it did here once. Sitting silent is what a server + does; the editor closing the pipe is what ends it, and that arrives as EOF + rather than as a timeout. + + Nothing in this agent puts a clock on the editor, in fact -- see + `ACPSession.call`. The two ends wait for each other indefinitely and either may + walk away, which is the arrangement the protocol actually describes. + """ + channel = os.fdopen(os.dup(1), "w", buffering=1) + os.dup2(2, 1) + sys.stdout = channel + try: + reader, writer = await acp.stdio.stdio_streams() + finally: + sys.stdout = sys.stderr + + # `run_agent`'s parameters are named from the client's point of view: the stream + # the client reads is the one this agent writes. + await acp.run_agent( + self, + input_stream=writer, + output_stream=reader, + use_unstable_protocol=True, + ) 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..ad61cd00b --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing.py @@ -0,0 +1,276 @@ +"""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. + + * ``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 `TenacityRetryer` 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). + +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 `TenacityRetryer` 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": ...})`` +""" + +import argparse +import asyncio +import collections.abc +import dataclasses +import enum +import pprint + +from auditing_agents import ( + Comparator, + Comparison, + Informalization, + Informalizer, + Strength, + Verdict, +) +from auditing_corpora import DOMAINS, Claim, Domain +from auditing_naive import NaiveAuditor + + +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. + """ + + domain: str + claim: Claim + statement: str + verdict: Verdict + explanation: str + comparison: Comparison | None = None + back_translation: Informalization | 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 = domain.statement_of(claim.theorem) + if mode is Mode.NAIVE: + judgement = NaiveAuditor().audit(claim.requirement, statement) + match, explanation = judgement.match, judgement.explanation + verdict = Verdict.CONFIRMED if match else Verdict.DISPUTED + return Audit(domain.name, claim, statement, verdict, explanation) + else: + back = Informalizer().informalize(statement) + comparison = Comparator().compare(claim.requirement, statement, back) + match, explanation = comparison.match, comparison.explanation + verdict = Verdict.CONFIRMED if match else Verdict.DISPUTED + return Audit( + domain.name, claim, statement, verdict, explanation, comparison, back + ) + + +async def audit_all( + domains: collections.abc.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 + ) + ) + ) + + +def pre_checks(audits: collections.abc.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 + + +def report(audits: collections.abc.Sequence[Audit], mode: Mode) -> None: + print(f"\n{'=' * 78}\nClaimCheck audit -- strategy: {mode.value}\n{'=' * 78}\n") + + for audit in audits: + pprint.pprint(audit) + + if notes := pre_checks(audits): + print("Pre-check diagnostics (deterministic, no model involved):") + for note in notes: + print(f" - {note}") + + 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) + + 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 "") + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + 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 or args.verify: + if all(domain.verify_corpus for domain in domains): + print( + f"VERIFIED: {sum(len(d.theorems) for d in domains)} theorems, 0 errors, no " + "`sorry`. Every claim below is proved.\nThe audit that follows is not " + "about whether they are true.\n" + ) + + if not args.verify_only: + 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..b6a4b892a --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing_agents.py @@ -0,0 +1,201 @@ +"""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. +""" + +import dataclasses +import enum +import typing + +import pydantic.dataclasses + +from effectful.handlers.llm import Skill + + +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' when and only when the theorem expresses the requirement. + Others are 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. + """ + + NONE = "none" + TAUTOLOGY = "tautology" + WEAKENED_CONCLUSION = "weakened-conclusion" + NARROWED_SCOPE = "narrowed-scope" + MISSING_CASE = "missing-case" + WRONG_PROPERTY = "wrong-property" + + +class Strength(enum.StrEnum): + """ + '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. + """ + + 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 + confidence: typing.Annotated[float, pydantic.Field(ge=0, le=1)] + + +@pydantic.dataclasses.dataclass(frozen=True) +class Comparison: + """Pass 2's verdict on one requirement/theorem pair. + + ``match`` is 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. + + ``__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. + """ + + match: bool + weakening: Weakening + discrepancy: str = dataclasses.field( + metadata={ + "description": "What the requirement asks for that the theorem does " + "not deliver. Empty when match is true." + } + ) + explanation: str + + 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 + + +class Informalizer: + """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.""" + + @Skill.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. + """ + + +class Comparator: + """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.""" + + @Skill.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} + + 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_corpora.py b/docs/source/llm_examples/autoformalization/auditing_corpora.py new file mode 100644 index 000000000..b490560bc --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing_corpora.py @@ -0,0 +1,840 @@ +"""The ClaimCheck benchmark: five Lean corpora, the claims made about them, and +the labelled ground truth the audit in `auditing.py` is scored against. + +Nothing in this module may reach an auditing agent. The harness builds a skill's +system prompt partly from the source of the module the skill is *defined* in, so +the unit of exposure is the module: the agents live in `auditing_agents` and +`auditing_naive`, and neither imports this file. The dependency runs the other +way -- this module imports `Verdict` from `auditing_agents` to spell the answer +key -- which is safe in exactly the direction that matters. See the module +docstring of `auditing_agents` for what happened when the corpora did share a +file with them. + +Split out of `auditing.py`, which drives the audit, scores it, and holds the +findings; here is only the data it runs on. +""" + +import dataclasses +import functools +import pathlib +import re +import sys +import textwrap + +from auditing_agents import Verdict + +# --------------------------------------------------------------------------- +# 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 +""" + + +# --------------------------------------------------------------------------- +# 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 skill, 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, ...] + + @functools.cached_property + def theorems(self) -> list[str]: + """The names of the theorems this corpus proves, in source order.""" + return re.findall(r"^theorem (\w+)", self.corpus, re.MULTILINE) + + def statement_of(self, qualified: str) -> str: + """Extract the *statement* of ``.`` from this corpus: + 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 = self.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 + + @functools.cached_property + def verify_corpus(self) -> bool: + """Compile this corpus with Lean, and say whether it came out proved. + + 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 this reuses its kernel rather than restating it, + imported inside the property as `world_model_agent.py` imports + `gridworlds`, so the example carries no Lean dependency unless the check + is asked for. + + Cached because importing Mathlib is by far the slowest thing this example + does, and a domain may be asked to verify more than once in a process. + ``False`` means the toolchain is not built; a corpus that *fails* raises + instead, since an unproved theorem makes every verdict about it + meaningless. + """ + # The examples are importable as ``docs.source.llm_examples...`` from the + # repository root, which is on ``sys.path`` under the harness but not when + # `auditing.py` 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.\nBuild it once (see formalization.py " + "--check-toolchain):\n" + " elan default stable\n" + f" cd {kernel.project} && lake exe cache get && lake build" + ) + return False + + print( + f"Compiling {self.name} ({len(self.theorems)} theorems) with " + "Lean 4 + Mathlib ..." + ) + result = kernel.compile(self.corpus) + if not result.ok: + raise SystemExit( + f"The {self.name} corpus does not compile:\n{result.messages}" + ) + if _SORRY.search(self.corpus): + raise SystemExit( + f"The {self.name} corpus contains `sorry`; it is not proved." + ) + return True + + +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), +} 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..9250106b6 --- /dev/null +++ b/docs/source/llm_examples/autoformalization/auditing_naive.py @@ -0,0 +1,55 @@ +import enum + +import pydantic.dataclasses + +from effectful.handlers.llm import Skill + + +class NaiveVerdict(enum.StrEnum): + """ + - **JUSTIFIED** if a theorem's statement expresses the requirement (it + may be stronger, that's fine). + - **NOT_JUSTIFIED** if there is a meaningful discrepancy: a theorem is + weaker, proves something different, is vacuous, or misses key aspects. + """ + + JUSTIFIED = "JUSTIFIED" + NOT_JUSTIFIED = "NOT_JUSTIFIED" + + +@pydantic.dataclasses.dataclass(frozen=True) +class NaiveJudgement: + """A verdict and a sentence of justification.""" + + verdict: NaiveVerdict + explanation: str + + @property + def match(self) -> bool: + return self.verdict is NaiveVerdict.JUSTIFIED + + +class NaiveAuditor: + """You check whether verified Lean theorems correctly formalize the natural + language requirements they are said to capture.""" + + @Skill.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 + + 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..bc661e498 --- /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 ``TenacityRetryer`` 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 + ``TenacityRetryer`` 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 ``TenacityRetryer`` 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 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 Skill, 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 ``TenacityRetryer`` 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 MRO) and is invisible to +# the closed-book NLProver, BlueprintAgent, and Reviewer. +# --------------------------------------------------------------------------- + + +class LeanAgent: + """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: + """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.""" + + @Skill.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``.""" + + @Skill.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: + """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.""" + + @Skill.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.""" + + @Skill.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: + """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.""" + + @Skill.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, ``TenacityRetryer``).""" + 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..116a6220f --- /dev/null +++ b/docs/source/llm_examples/autoresearch/illustration.py @@ -0,0 +1,685 @@ +"""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 ``TenacityRetryer`` (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 Skill 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 ``TenacityRetryer`` -- "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 Skills. + +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 + ``TenacityRetryer`` +- 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 +- 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 evaluation. The paper's VLM-as-a-Judge scores a render against a human-drawn +# figure on four dimensions and aggregates them hierarchically; scoring the pipeline +# is out of scope here, as it is in the sibling examples. ``refine`` still returns +# both ends of the paper's Critic-on/off ablation -- the round-0 and final renders -- +# and every round is written to ``outdir``, so the loop's effect can be read off the +# PNGs directly. + +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 Skill, 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 skill'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 inspects. 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: + """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()) + + @Skill.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: + """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.""" + + @Skill.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: + """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.""" + + @Skill.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} + """ + + @Skill.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: + """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.""" + + @Skill.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 Skill. +# --------------------------------------------------------------------------- + + +class Critic: + """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.""" + + @Skill.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). + + Runs the paper's *fixed* T rounds: the Critic always emits a refined + description P_{t+1}, and it is always re-rendered. Halting early on an empty + ``issues`` list would make the demonstration conditional on the Critic's + mood -- a model that reports no issues on I_0 (the common case on a simple + plot) would leave ``round_1.png`` unwritten and the refinement loop, the + thing this function exists to show, entirely unexercised. + + 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) + plan = critique.refined + img = render(Visualizer().visualize(plan), outdir / f"round_{t + 1}.png") + + 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). Returns the final plan; both rounds' + renders are left on disk under ``outdir``.""" + # Linear Planning Phase. The Retriever certifies its selection as *keys*, so + # resolve them against R here -- the Planner learns from the exemplars' + # structure and style notes, which a bare key does not carry. + retrieval = Retriever().retrieve(task) + exemplars = [REFERENCES[key] for key in retrieval.selected] + description = Planner().plan(task, exemplars) + + 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..3bdd9cc52 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/implementation.py @@ -0,0 +1,806 @@ +"""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 + ``TenacityRetryer`` (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 ``TenacityRetryer``, + 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 Skill + +# 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 +# skill'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 `write_and_run_body` tool) and the deterministic evaluator. +# This is the distinctive shape of a coding agent, versus the literature examples' +# search tools. +# --------------------------------------------------------------------------- + + +class Designer: + """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.""" + + @Skill.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: + """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.""" + + @Skill.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: + """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.""" + + @Skill.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) + # Tie-break away from `best`: `max` and `min` both return the *first* + # element among equals, so on a run where every rollout scores the same + # -- common early, before the reward spreads out -- `worst` would be the + # very object `best` is, the guard below would fire, and comparative + # reflection would be skipped in silence for the whole search. + worst = min(ok, key=lambda r: (r.reward, r.path == best.path)) + 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..d5350d74b --- /dev/null +++ b/docs/source/llm_examples/autoresearch/investigation.py @@ -0,0 +1,884 @@ +"""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 ``TenacityRetryer`` 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 ``TenacityRetryer`` 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 ``Skill`` + 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 ``Skill`` 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 Skill, 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 skill +# 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: + """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]] + + @Skill.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: + """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.""" + + @Skill.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 + + +@Skill.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: + """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}") + + @Skill.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} + """ + + @Skill.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 + + +@Skill.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} + """ + + +@Skill.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} + """ + + +@Skill.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]]: + 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 ``TenacityRetryer``) 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 TenacityRetryer, forces a correction. The + # skill is told to cite a fabricated reference; certification bounces the + # first attempt and the model must ground it before the call can return. + @Skill.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 TenacityRetryer, 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..48710f942 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/review.py @@ -0,0 +1,643 @@ +"""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 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 Skill 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 + ``TenacityRetryer`` 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 class, offered to subclass skills via the + MRO but invisible to the sibling toolless agents -- tool scoping as + encapsulation, so no skill 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 Skill 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 Skill, 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 skill'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`` (below) delegates to whichever backend main() +# selects; both populate RETRIEVED, the registry a MissingBaseline is certified +# against, so grounding means "cite only work you actually retrieved" regardless +# of source. +# --------------------------------------------------------------------------- + +# Papers returned by ``search`` this run, keyed by citation key (a corpus key or a +# Semantic Scholar paperId). Shared across the concurrently-running Historian and +# Scout; plain key insertion is safe under the GIL. +RETRIEVED: dict[str, LitEntry] = {} + + +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." + } + ) + + def __post_init__(self) -> None: + if self.paper_key not in RETRIEVED: + raise ValueError( + f"paper_key {self.paper_key!r} was not among the papers search " + f"returned this run (retrieved: {sorted(RETRIEVED)}); cite only " + f"omitted work you actually retrieved via the search tool" + ) + + +@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: + """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.""" + + @Skill.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' skills and to no one else. +# --------------------------------------------------------------------------- + + +class Scholar: + """Base for agents that read the literature. The ``search`` tool defined here + is inherited (via the MRO) by every ``Scholar`` subclass's skills, + 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) + RETRIEVED.update(found) # certify later citations against what was retrieved + 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.""" + + @Skill.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.""" + + @Skill.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: + """You are the Question Generator. You turn the gathered context into a few + sharp, specific probing questions aimed at a submission's weakest points.""" + + @Skill.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: + """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.""" + + @Skill.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: + """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 + + @Skill.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.""" + RETRIEVED.clear() # each review is grounded only in what it retrieves + + # 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..e1f649e55 --- /dev/null +++ b/docs/source/llm_examples/autoresearch/writing.py @@ -0,0 +1,733 @@ +"""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 Skill 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 + ``TenacityRetryer`` 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`` Tool; the Outline, Plotting, Section, and + Refinement agents define no Tool at all 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 Skill 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 Skill -- 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 ``TenacityRetryer`` 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 +# ``Citation.__post_init__`` both read a tiny in-memory INDEX, so a retrieved paper +# has a stable key a Citation can certify against. The authentication half of the +# loop is therefore a decode-time check on the type rather than a second Tool call. +# 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 Skill, 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 skill'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 skill 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 Skill. +# --------------------------------------------------------------------------- + + +@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] + + +# Every field below that carries LaTeX carries this too. The harness answers a +# Skill by having the model write a Python function body, so a control sequence in +# an ordinary string literal is escape-processed before it is ever a value: +# ``"\texttt"`` is a TAB followed by ``exttt``, and ``"$3.1\times$"`` loses the +# ``\t`` the same way. The damage is silent -- the manuscript simply comes out with +# a tab in the middle of a word -- so the guidance goes on the fields themselves, +# where the model reads it as part of the schema it is filling. +_LATEX_ESCAPING = ( + "Contains LaTeX control sequences. When writing this as a Python string " + "literal, use a raw string (r'...') or double every backslash: in an ordinary " + r"literal ``\texttt`` and ``\times`` collapse to a TAB character." +) + + +# --------------------------------------------------------------------------- +# 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 web_search" + ) + 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 = dataclasses.field(metadata={"description": _LATEX_ESCAPING}) + latex: str = dataclasses.field( + metadata={ + "description": "Self-contained LaTeX for the figure: a pgfplots axis or " + "tabular for a plot, TikZ for a diagram. " + _LATEX_ESCAPING + } + ) + + +@pydantic.dataclasses.dataclass(frozen=True) +class Section: + """One body section of the manuscript (Method, Experiments, ...).""" + + name: str + body: str = dataclasses.field(metadata={"description": _LATEX_ESCAPING}) + + +@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: + """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.""" + + @Skill.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: + """You are the Literature Review Agent. You run a two-move discovery loop -- + *identify* candidate prior work with web search, and *verify* it by citing it: + every ``Citation`` you build is authenticated as it is constructed, and one that + names no indexed paper, or one published on or after the cutoff, is rejected and + returned to you. You 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 + + @Skill.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. 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: + """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.""" + + @Skill.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: + """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.""" + + @Skill.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: + """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 Skill: a + module-level ``@Skill.define`` lands in every other skill's lexical scope + and is offered to those agents as a callable tool, but a Skill *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 + + @Skill.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: + """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.""" + + @Skill.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, related = await asyncio.gather( + asyncio.to_thread( + PlottingAgent().draw, outline.figures, materials.experimental_log + ), + 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..b653093a4 --- /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 Skill + + +@dataclasses.dataclass +class ChatBot: + """Conversational agent that remembers the conversation so far.""" + + bot_name: str + + @Skill.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..5baea479a --- /dev/null +++ b/docs/source/llm_examples/basics/error_recovery.py @@ -0,0 +1,96 @@ +"""Recovering from failed LLM output: flaky tools and invalid structured output. + +A single task -- rate a movie after looking it up -- drives the tool retry path: + +Demonstrates: +- TenacityRetryer surfacing tool exceptions back to the LLM as tool messages, so a + flaky tool (lookup_movie) can succeed after multiple attempts +- the same feedback loop standing behind structured output: a `Rating` whose + ``__post_init__`` rejects it comes back as a pydantic validation error for the + model to correct + +The second path is a guard, not a demonstration: the harness puts this module's +whole source in the system prompt, ``__post_init__`` included, so the model can +read the rule it must satisfy and usually gets a valid `Rating` first try. That +is the guard working, not the retry failing to fire -- to watch it fire, tighten +the rule to something the source does not spell out. +""" + +import argparse +import dataclasses +import typing + +from effectful.handlers.llm import Skill, 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}'" + ) + + +# --------------------------------------------------------------------------- +# Skill: uses the flaky tool, returns validated structured output +# --------------------------------------------------------------------------- + + +@Skill.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..e02f72f7a --- /dev/null +++ b/docs/source/llm_examples/basics/flight_booking.py @@ -0,0 +1,280 @@ +"""Flight booking: composing agents by passing one's typed output to the next. + +Demonstrates: +- A standalone ``@Skill.define`` (``extract_flights``) whose typed result is + stored as an ``Agent`` field and spliced into a second agent's prompt +- A post-condition on a skill's return type: a ``pydantic.AfterValidator`` + that compares the answer against the arguments the call was made with, so + rejecting a wrong answer and retrying is the harness's job +- A pre-condition on a skill's parameter, guarding it with a second skill's + judgement: an off-topic seat request is rejected before any booking happens +- Interactive human-in-the-loop flow +- ``Agent`` history for conversational seat selection +""" + +import argparse +import dataclasses +import datetime +import enum +from typing import Annotated + +import annotated_types +import pydantic + +from effectful.handlers.llm import Skill + +# --------------------------------------------------------------------------- +# 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: Annotated[int, pydantic.Field(gt=0)] + origin: Airport # three-letter airport code + destination: Airport # three-letter airport code + date: datetime.date # YYYY-MM-DD + + +class Seat(enum.StrEnum): + """Seats A and F are window seats. Seats C and D are aisle seats.""" + + A = "A" + B = "B" + C = "C" + D = "D" + E = "E" + F = "F" + + +@pydantic.dataclasses.dataclass(frozen=True) +class SeatPreference: + """ + User's seat preference extracted from natural language. + + Row 1 is the front row with extra legroom. + Rows 14 and 20 also have extra legroom. + """ + + row: Annotated[int, pydantic.Field(ge=1, le=30)] + seat: Seat + + +# --------------------------------------------------------------------------- +# 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 skill (inner "agent") +# --------------------------------------------------------------------------- + + +@Skill.define +def extract_flights(web_page_text: str) -> list[FlightDetails]: + """Extract all flight details from the following text. + + {web_page_text} + """ + + +# --------------------------------------------------------------------------- +# Post-condition on the search result (plain Python, no LLM needed) +# --------------------------------------------------------------------------- + + +def matches_request( + flight: FlightDetails, info: pydantic.ValidationInfo +) -> FlightDetails: + """Check that the flight the model chose matches the criteria it was asked for. + + A skill's arguments are the validation context its answer is decoded under, + so a post-condition can compare that answer against the request without + being closed over it: ``info.context`` holds this call's ``origin``, + ``destination`` and ``date`` (and ``self``), whichever way the model + answered. Raising rejects the answer -- the harness feeds the message back + and the model tries again, up to ``--num-retries`` times, after which the + call raises rather than returning a flight nobody asked for. + """ + request = info.context or {} + errors = [ + f"{field} should be {request[field]}, got {getattr(flight, field)}" + for field in ("origin", "destination", "date") + if getattr(flight, field) != request[field] + ] + if errors: + raise ValueError("; ".join(errors)) + return flight + + +# --------------------------------------------------------------------------- +# Flight search agent +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class FlightFinder: + """Agent that finds flights matching user criteria.""" + + available_flights: list[FlightDetails] + + @Skill.define + def find_flight( + self, origin: Airport, destination: Airport, date: datetime.date + ) -> Annotated[FlightDetails, pydantic.AfterValidator(matches_request)]: + """ + Find the cheapest flight from {origin} to {destination} on {date}. + + List of available flights (from the web page): + {self.available_flights} + """ + + +@Skill.define +def is_seat_request(user_input: str) -> bool: + """ + Determine whether the user's message is about where they want to sit on + the plane -- a seat, a row, or a seating preference: {user_input} + Do not use any tools. + """ + + +class SeatSelector: + """Agent that extracts seat preferences from natural language.""" + + @Skill.define + def select_seat( + self, user_input: Annotated[str, annotated_types.Predicate(is_seat_request)] + ) -> SeatPreference: + """Extract the user's seat preference from their message. + + {user_input} + """ + + +# --------------------------------------------------------------------------- +# Booking flow +# --------------------------------------------------------------------------- + + +def book_flight( + origin: Airport, + destination: Airport, + date: datetime.date, + interactive: bool = False, +) -> None: + """End-to-end flight booking with search, validation, and seat selection.""" + searcher = FlightFinder(available_flights=extract_flights(FLIGHTS_PAGE)) + + # --- Search (checked by `find_flight`'s post-condition, so an answer that + # doesn't match the request is rejected and retried before it reaches here) --- + flight = searcher.find_flight(origin, destination, date) + + 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: + try: + seat = selector.select_seat(request) + print(f" Seat: row {seat.row}, seat {seat.seat}") + except pydantic.ValidationError: + # The pre-condition rejected the message, so the model was never + # asked to read a seat out of it. The predicate reports only that it + # failed, so the message a person sees is the caller's to write. + print(f" Rejected: {request!r} is not a seat preference.") + return + + 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..56ad82518 --- /dev/null +++ b/docs/source/llm_examples/basics/guardrails.py @@ -0,0 +1,67 @@ +"""Travel advisor with input guardrails. + +Demonstrates: +- A pre-condition on a skill's parameter: one skill's judgement, attached to + the annotation as ``annotated_types.Predicate`` metadata, guarding another + skill's input wherever that input comes from +- A post-condition on the return: an ordinary Python predicate, attached the + same way, enforced by the decoder that reads the model's answer -- a + rejection is fed back to the model, which then answers again +""" + +import argparse +import typing + +import annotated_types +import pydantic + +from effectful.handlers.llm import Skill + + +@Skill.define +def is_safe_query(user_query: str) -> bool: + """ + Determine whether the user's query is purely related to travel advice: {user_query} + """ + + +def is_concise_answer(answer: str) -> bool: + """Determine whether the answer is concise (<100 words).""" + return len(answer.split()) < 100 + + +@Skill.define +def travel_query( + user_query: typing.Annotated[str, annotated_types.Predicate(is_safe_query)], +) -> typing.Annotated[str, annotated_types.Predicate(is_concise_answer)]: + """ + Produce a concise (<100 word) answer to: {user_query} + """ + + +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(f"Query: {query}") + try: + print("Answer:", travel_query(query)) + except pydantic.ValidationError: + # The guard reports only that its predicate failed, so the message + # a person sees is the caller's to write. + print(f"Rejected: '{query}' is not related to travel advice.") + + +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..29bafba03 --- /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 Skill, 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 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 + + @Skill.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..d0cbb341b --- /dev/null +++ b/docs/source/llm_examples/basics/image_input.py @@ -0,0 +1,50 @@ +"""Passing PIL images directly to a skill. + +Demonstrates: +- Skills 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 Skill + + +@Skill.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..02c3b62d9 --- /dev/null +++ b/docs/source/llm_examples/basics/image_tool.py @@ -0,0 +1,114 @@ +"""Giving an agent tools that operate on images. + +Demonstrates: +- An ``Agent`` subclass whose ``Tool`` methods manipulate ``PIL.Image.Image`` values +- Passing opaque integer handles across the model boundary so images stay in Python +- A ``Skill`` that plans a sequence of tool calls to build a composite image +""" + +import argparse +import pathlib + +from PIL import Image + +from effectful.handlers.llm import Skill, Tool + + +class ImageTools: + """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) + + @Skill.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__) + parser.add_argument( + "--image", + type=str, + default=str(DEFAULT_IMAGE), + metavar="PATH", + help="Path to the input image to rotate-and-concatenate.", + ) + parser.add_argument( + "--interactive", + action="store_true", + help="Display the resulting image in an image viewer.", + ) + args = parser.parse_args() + + image_agent = ImageTools() + img = Image.open(args.image) + + result = image_agent.rotate_and_concat(img) + print(f"Result image: {result.width}x{result.height}") + + if args.interactive: + result.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..697e3f4ad --- /dev/null +++ b/docs/source/llm_examples/basics/lexical_scope.py @@ -0,0 +1,99 @@ +"""Composition via lexical scope: auto-captured sub-skills, invoked two ways. + +Demonstrates: +- Module-level @Skill.define sub-skills auto-captured into other skills' + lexical scope, with no explicit registration +- An Agent grouping @Tool.define tools with a @Skill.define orchestrator that + calls those tools and the sub-skills directly (model-driven composition) +- A skill returning a Callable: the model synthesizes a function that calls the + same sub-skills 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 Skill, Tool + + +@Skill.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.""" + + +@Skill.define +def story_funny(topic: str) -> str: + """Write a funny, humorous story about {topic}. Do not use any tools.""" + + +class TripPlanner: + """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") + + @Skill.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 skill calls tools and sub-skills. + print("=== Orchestrator skill (model-driven composition) ===") + planner = TripPlanner() + print(planner.plan_trip_story(args.style)) + + elif args.method == "code": + + @Skill.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-skill for humor, or `story_with_moral` for a lesson.""" + + # (2) Code-driven: the model synthesizes a function that calls the sub-skills. + 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..bef0e6f22 --- /dev/null +++ b/docs/source/llm_examples/basics/map_reduce.py @@ -0,0 +1,140 @@ +"""Map-reduce resume evaluation. + +Demonstrates: +- Fan-out: evaluating multiple items independently with the same skill +- 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 + +import annotated_types + +from effectful.handlers.llm import Skill + +# --------------------------------------------------------------------------- +# Structured output +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class Evaluation: + name: str + qualified: bool + strengths: str + weaknesses: str + score: typing.Annotated[int, annotated_types.Ge(1), annotated_types.Le(10)] + + +# --------------------------------------------------------------------------- +# Skills +# --------------------------------------------------------------------------- + + +@Skill.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} + """ + + +@Skill.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 skill 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..d713bc795 --- /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 (``@Skill.define``) +""" + +import argparse +import dataclasses + +import litellm +import numpy as np + +from effectful.handlers.llm import Skill, 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` skill +# share one instance, so the tool is auto-captured from lexical scope. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class RAGAgent: + """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) + + @Skill.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="text-embedding-3-small", + 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..3040fa132 --- /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 skills from lexical scope +- An Agent subclass with persistent conversation history +- One Skill 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 Skill, 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 that answers research questions using web search, refining on feedback.""" + + @Skill.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}""" + + @Skill.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) +# --------------------------------------------------------------------------- + + +@Skill.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..b9c4c7b81 --- /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 ``@Skill.define`` +- Executing SQL against a real SQLite database +- Feeding execution errors back to the LLM for iterative fixing +- Reading the live schema out of SQLite and splicing it into the prompt +""" + +import argparse +import sqlite3 +import textwrap + +from effectful.handlers.llm import Skill + +# --------------------------------------------------------------------------- +# 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]) + + +# --------------------------------------------------------------------------- +# Skills +# --------------------------------------------------------------------------- + + +@Skill.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. + """ + + +@Skill.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..e4164fbe2 --- /dev/null +++ b/docs/source/llm_examples/choreographies/library.py @@ -0,0 +1,619 @@ +"""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 skill 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.types.Agent` an ``agent_id`` and install +`~effectful.handlers.llm.harness.durability.persistence.SQLitePersister` to checkpoint it.) + +## Why the program is async, and where the threads went + +`effectful`'s handler stack is synchronous, top to bottom: `call_agent`, +`fwd`, and everything in `effectful.handlers.llm.harness` 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 skill 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 skill 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 + ) + reviews = choreo( + "Build a URL slugify library", + architect=architect, + coder=[coder1, coder2], + reviewer=reviewer, + ) + +Calling it is all there is to it -- the skill calls inside need a model behind +them, which the module launcher supplies:: + + python -m effectful.handlers.llm.harness your_choreography.py --model gpt-4o-mini + +``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]( + skill: 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 *skill* 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 + skill 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 skill 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(skill, *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 skill 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 skill 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._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 + + @property + def _agent_id(self) -> str: + return self._agent.__agent_id__ + + @implements(step) + def _step(self, skill: Callable, *args, **kwargs) -> Awaitable: + return self._run_step(self._next_step(), skill, 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, skill: 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. + """ + if ( + hasattr(skill, "__history__") + and (agent_id := skill.__self__.__agent_id__) != self._agent_id # type: ignore + ): + raise RuntimeError( + f"a scatter item taken on by {self._agent_id!r} called " + f"{skill.__name__}(), which belongs to " + f"{agent_id!r}. A scatter item is work one agent " + f"does alone; step across agents outside the scatter." + ) + return self._in_thread(skill, *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, skill: Callable, args: tuple, kwargs: dict + ) -> Any: + + if not hasattr(skill, "__history__"): + # Unbound skill: not owned by anyone, so every agent runs it. + return await self._in_thread(skill, *args, **kwargs) + + agent: Agent = skill.__self__ # type: ignore + if self._agent_ids is not None and agent.__agent_id__ not in self._agent_ids: + raise ChoreographyError( + f"{skill.__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(skill, *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 skill + 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.harness.durability.persistence.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]) + + result = choreo( + "Build a library...", + architect=architect, + coder=coder, + reviewer=reviewer, + ) + + Run the calling script under ``python -m effectful.handlers.llm.harness`` to + put a model behind the skill calls. + """ + + 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..2209c928c --- /dev/null +++ b/docs/source/llm_examples/choreographies/multi_agent_choreography.py @@ -0,0 +1,379 @@ +"""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 Skill, 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: + """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, *, __agent_id__: str): + self.__agent_id__ = __agent_id__ + 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) + + @Skill.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: + """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, *, __agent_id__: str): + self.__agent_id__ = __agent_id__ + 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}" + + @Skill.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: + """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, + *, + __agent_id__: str, + ): + self.__agent_id__ = __agent_id__ + 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:]}" + + @Skill.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/avo.py b/docs/source/llm_examples/optimization/avo.py new file mode 100644 index 000000000..844794fcf --- /dev/null +++ b/docs/source/llm_examples/optimization/avo.py @@ -0,0 +1,451 @@ +"""Agentic variation: the evaluator moves inside the model's reach (AVO). + +Implements "AVO: Agentic Variation Operators for Autonomous Evolutionary Search" +(arXiv 2603.24517). Its argument is against the system the rest of this directory +implements. Evolutionary search with an LLM decomposes variation as +``Vary(P) = Generate(Sample(P))``, and the model is confined to ``Generate``: it emits +one candidate per invocation, having tested nothing. AVO replaces the entire operator +with a single autonomous agent run, ``Vary(P) = Agent(P, K, f)``, where the agent is +handed the lineage ``P``, a knowledge base ``K``, and -- the load-bearing part -- the +scoring function ``f`` itself, as a utility it may call whenever it likes. It edits, +evaluates, reads the diagnostics, repairs, and only then commits. + +The whole of that change, in this library, is four lines of ordinary Python: + + 1. **``f`` becomes a `Tool`** in the variation skill's lexical scope. The multi-turn + agent loop the paper builds around it is what a `Skill` call already is -- the + model calls `VariationAgent.evaluate` as many times as its budget allows before + answering -- so the paper's central mechanism is a scope change, not a subsystem. + 2. **The lineage is in reach**: the incumbent arrives as a typed ``Kernel`` + argument, rendered into the prompt as its own source, and stays callable *by + name* in the REPL, so the agent can time it against its candidate. + 3. **One persistent agent** spans the whole run, where `kernels.py` builds a fresh + `Proposer` per iteration. Its history is the paper's memory, and is what the + transfer episode below rides on. + 4. **The commit rule inverts**: the agent self-certifies, and the framework shrinks + to `evolve_lineage` -- verify what came back, commit if it earns a place. + +The domain is `kernels.py`'s, deliberately: both examples score through one +`evaluate_kernel`, so the difference between them is the operator and not the +yardstick. The artifact here is the kernel *itself*, which is the paper's artifact, +where `kernels.py` evolves the instruction that writes one. + +Demonstrates: +- An evaluator handed to the model as a tool, and metered, so "how many times did the + agent score something" is a number in the report rather than an unknown +- A synthesized callable passed *by reference* to a tool: the agent defines a kernel in + the REPL and writes ``self.evaluate(candidate)``, which the harness type-checks as a + real Python call expression -- no serialization of the artifact anywhere +- The stdlib as the knowledge base: ``help(math)`` in the REPL is the paper's + documentation-retrieval channel, with no task-specific retrieval tool +- `TenacityRetryer` plus decode-time doctests as the paper's implicit repair stage +- An EVO control (``--evo``) that is the *same domain* through `optimize_anything`'s + single-turn proposer, reported in both currencies -- though see below for why this + domain cannot make that a test of anything +- Transfer (``--transfer``): the same agent, with its accumulated history, adapting its + kernel to a sibling task -- the paper's 30-minute MHA -> GQA episode +""" + +# Simplifications vs. the source: +# - The domain cannot discriminate the two operators, which is the load-bearing +# simplification and the one to read first (see "What this example does not do"). The +# task was chosen for having a measurable optimization ladder and a correctness cliff; +# it should also have been chosen for having a ladder whose *order a model cannot +# predict*, since that is the only condition under which holding the evaluator can +# beat guessing. Every other entry in this list is a simplification of the paper's +# setup; this one is a limitation of the experiment. +# - Pure-Python list transforms on a CPU, not CUDA kernels on a B200 against cuDNN and +# FlashAttention-4. The baseline is a deliberately plain reference implementation +# rather than months of expert tuning, so "3x the baseline" here and "3.5% over +# cuDNN" there are not the same kind of quantity, and the gap between them is most +# of what makes the paper's result hard -- and most of why an agent that can profile +# is worth its cost there and is not here. +# - Minutes and a handful of steps against 7 days and 40 committed versions. The +# paper's trajectory -- discrete jumps separated by plateaus -- needs a budget this +# example does not spend; expect one large jump onto the best rung of the ladder and +# then grinding, which is what the ladder above predicts. +# - No supervisor. The paper adds one that watches for stagnation and steers the search +# when it plateaus, but never ablates it: there are no intervention counts and no +# with/without arm, so its contribution is asserted rather than measured. At this +# budget a stagnation trigger would likely never fire. Restoring it is a stagnation +# counter plus one `redirect` skill on a *separate* agent -- separate so that it +# reviews from outside the context that plateaued, and so that lexical scope keeps +# `evaluate` out of its reach. +# - The knowledge base is the standard library rather than PTX ISA documentation, and +# the agent reaches it the same way the paper's does: by reading docs in its own +# session, not through a retrieval tool written for the task. +# - No git. The lineage is a list in memory; the paper commits each version, which is +# how its run survives a restart. `Agent.__agent_id__` plus ``--persist-db`` would be +# the equivalent here and is not wired up. +# - The two arms cannot be matched on both currencies, and this is not a defect of the +# comparison but of what the two operators buy. One AVO step spends one proposal and +# as many evaluator calls as its budget allows; one EVO iteration spends one of each. +# ``--evo`` matches on evaluator calls, which is the currency the paper's claim is +# about (the agent's advantage is supposed to be worth what it costs to test), and +# the report prints both counts for both arms so a reader can see the other side of +# the trade. +# - One run per arm and no variance estimate, against a timing ratio whose noise floor +# is a few percent (visible in the seed's own score, which is 1.0 by construction and +# does not measure as exactly 1.0). Differences smaller than that are noise, and the +# report says so rather than ranking the arms. +# - Both arms see this module's source, because the harness puts it in the system +# prompt. The test cases and the reference implementation are therefore *not* hidden +# from either arm, whatever `kernels.py` says about its own: the enforceable +# asymmetry between the arms is that the AVO agent is *offered* a metered evaluator +# and the EVO proposer is not, and a control that chose to reconstruct scoring for +# itself in the REPL would be neither prevented nor undetected -- the report prints +# both arms' evaluator counts partly so that this stays visible. + +import argparse +import collections.abc +import random +import time + +from docs.source.llm_examples.optimization.kernels import ( + KERNEL_TASKS, + Kernel, + KernelTask, + evaluate_kernel, +) +from docs.source.llm_examples.optimization.library import ( + Diagnostic, + Evaluation, + Lineage, + Result, + Rollout, + evolve_lineage, + optimize_anything, + source_of, +) +from effectful.handlers.llm import Skill, Tool + +# The timed configurations: the paper scores each kernel on a vector of benchmark +# shapes, and keeping several here is what stops a candidate from winning by being +# good at one size. Small enough that a whole evaluation costs a fraction of a second, +# because the agent runs many of them per step. +CONFIGS: tuple[int, ...] = (30_000, 100_000, 300_000) + +TASKS: dict[str, KernelTask] = {task.name: task for task in KERNEL_TASKS} + +# "MHA" and "GQA": the task evolved, and the sibling the result is transferred to. +TASK = TASKS["zscore"] +TRANSFER_TASK = TASKS["l2_normalize"] + + +class VariationAgent: + """You are a performance engineer optimizing one small Python kernel, over a long + run in which you will be asked for improvements repeatedly. + + You have a real evaluator: `evaluate` runs the kernel you hand it against the + hidden correctness cases and then times it against the reference implementation on + several input sizes. USE IT. Write a candidate in the REPL, evaluate it, read what + came back, and fix what it tells you -- the whole point of the loop you are in is + that you can measure instead of guessing. An idea you have not evaluated is not an + improvement. + + Two things are worth remembering across the whole run, because you are the same + agent each time: which optimizations actually paid, and which looked promising and + did not. Say so when you answer, and if the transcript gets long, compact it. + """ + + def __init__( + self, + task: KernelTask, + configs: collections.abc.Sequence[int] = CONFIGS, + step_budget: int = 6, + ): + self.task = task + self.configs = tuple(configs) + self.step_budget = step_budget + self.spent = 0 # evaluator calls in the current variation step + self.total = 0 # ... and over the whole run, which is the reported currency + + @Tool.define + def evaluate(self, kernel: Kernel) -> Evaluation: + """Score a candidate kernel: correctness first, then speed. + + The kernel is run against hidden correctness cases -- including degenerate inputs + the specification mentions -- and then timed on several large inputs against the + reference implementation. The score is the geometric mean of the per-size + speedups, and an incorrect kernel scores zero however fast it is. + + Pass the function itself, not its source: define it with the REPL tool and then + call this one on the name you bound. + """ + self.total += 1 + if self.spent >= self.step_budget: + return Evaluation( + score=0.0, + diagnostics=[ + Diagnostic( + "BUDGET EXHAUSTED", + f"This is NOT a score for your kernel -- it is a refusal to " + f"measure. You have used all {self.step_budget} evaluations " + f"for this step. Return the best kernel you have already " + f"measured.", + ) + ], + ) + self.spent += 1 + evaluation = evaluate_kernel(kernel, self.task, self.configs) + return Evaluation( + score=evaluation.score, + metrics=evaluation.metrics, + diagnostics=[ + *evaluation.diagnostics, + Diagnostic( + "evaluations remaining", + f"{self.step_budget - self.spent} of {self.step_budget} left in " + f"this step", + ), + ], + ) + + @Skill.define + def vary(self, current: Kernel, evaluation: Evaluation) -> Kernel: + """Improve this kernel. + + + {self.task.spec} + + + The current best kernel is below, and it is also bound in your REPL session as + ``current``, so you can time your candidate against it directly. + + + {current} + + + Here is how it scored. Read it before writing anything: it reports the speedup + at each input size, which failing cases there were, and the code that was + measured. + + + {evaluation} + + + Work like an engineer, not like an oracle: + + - Write candidates in the REPL and hand each one to `evaluate`. You have a + budget of evaluations per step; the diagnostics tell you how many are left. + - When a candidate is slower or wrong, the diagnostics say which case failed + and how fast it ran. Diagnose it before you try again. + - Consult the standard library rather than guessing at it: ``help(math)``, + ``import numpy`` and the like all work in the REPL, and whether a function + exists and what it costs are both things you can check. + - Watch the degenerate cases. The specification names them, and the fastest + arithmetic is often exactly the arithmetic that gets them wrong. + + Return a kernel that you have MEASURED to be correct and at least as fast as + the one you were given. It is re-scored after you return it, and it is rejected + if it does not hold up so returning something you did not evaluate wastes the step. + + Your function's docstring MUST contain doctests certifying its contract, and + they are run before your answer is accepted. Write at least one that checks the + degenerate input the specification describes, 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). + """ + + +def run_avo(args: argparse.Namespace) -> tuple[Lineage[Kernel], VariationAgent]: + """Agentic variation: one persistent agent, one lineage, the evaluator in reach.""" + agent = VariationAgent(TASK, CONFIGS, args.step_budget) + + def vary(current: Kernel, evaluation: Evaluation) -> Kernel: + # The per-step budget resets here rather than inside the tool: the tool cannot + # see where one variation step ends and the next begins, and the agent's + # history spans all of them. + agent.spent = 0 + return agent.vary(current, evaluation) + + lineage = evolve_lineage( + vary=vary, + evaluator=lambda kernel: evaluate_kernel(kernel, TASK, CONFIGS), + seed=seed_kernel(TASK), + budget=args.budget, + ) + return lineage, agent + + +def run_evo(args: argparse.Namespace, rng: random.Random) -> Result: + """The control: `optimize_anything` in single-task mode over the same artifact. + + Single-task mode, not a one-element dataset: with no dataset the Pareto objectives + are the evaluation's own sub-scores, which here are the per-configuration speedups + -- so the frontier can hold a candidate that wins only at one input size, which is + the mode's whole purpose. + """ + + @Skill.define + def propose_kernel(current: Kernel, feedback: list[Rollout]) -> Kernel: + """You are a reflective optimizer. Rewrite this kernel to be faster. + + + {current} + + + Here is how it scored, including the speedup at each input size and any failing + cases: + + + {feedback} + + + Diagnose what is costing the most, then write a faster kernel that computes exactly + the same thing -- including on the degenerate inputs, since an incorrect kernel + scores zero however fast it is. + + Your function's docstring MUST contain doctests certifying its contract, prefixing + each input line with the doctest prompt (three ``>`` characters and a space). + """ + + return optimize_anything( + evaluator=lambda kernel, _: evaluate_kernel(kernel, TASK, CONFIGS), + proposer=propose_kernel, + seed=seed_kernel(TASK), + budget=args.evo_budget or args.budget * args.step_budget, + selection=args.selection, + rng=rng, + task_name=TASK.name, + ) + + +def seed_kernel(task: KernelTask) -> Kernel: + """Candidate zero: the reference implementation itself. + + Both arms start from it, and it is also the denominator of the score, so the seed + measures 1.0 by construction -- give or take the timing noise, which is worth + seeing. "The search improved on its seed" therefore means "it beat the + straightforward implementation", with no baseline offset to argue about. + """ + from docs.source.llm_examples.optimization.kernels import REFERENCE + + return REFERENCE[task.name] + + +def transfer(agent: VariationAgent, kernel: Kernel) -> tuple[Evaluation, Evaluation]: + """The paper's MHA -> GQA episode: the same agent, a sibling task, one step. + + Nothing is reset. The agent keeps the history in which it discovered whatever it + discovered about ``zscore``, and is asked for a kernel for a task that shares that + task's shape but not its answer. The paper reports 30 minutes of autonomous + adaptation for the same move. + + Returns the sibling task's seed evaluation and the adapted kernel's, so the report + can state the transfer as a ratio against the reference rather than as a claim. + """ + agent.task, agent.spent = TRANSFER_TASK, 0 + before = evaluate_kernel(seed_kernel(TRANSFER_TASK), TRANSFER_TASK, CONFIGS) + adapted = agent.vary(kernel, before) + return before, evaluate_kernel(adapted, TRANSFER_TASK, CONFIGS) + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- + + +def speedups_of(evaluation: Evaluation) -> str: + return ", ".join( + f"{m.name.removeprefix('speedup@')}: {m.value:.2f}x" + for m in evaluation.metrics + if m.name.startswith("speedup@") + ) + + +def report_avo(lineage: Lineage[Kernel], agent: VariationAgent, seconds: float) -> None: + """The trajectory, the two currencies, and the winning kernel.""" + print("\n" + "=" * 72) + print( + f"mode: agentic variation (AVO) | {lineage.attempts} variation steps in " + f"{seconds:.0f}s" + ) + print( + f"seed {lineage.seed.score:.6g} -> best {lineage.best.score:.6g} " + f"({len(lineage.versions) - 1} committed, {lineage.rejected} rejected)" + ) + print( + f"cost: {agent.total} evaluator calls by the agent + {lineage.evaluations} " + f"verifications by the framework = {agent.total + lineage.evaluations} total, " + f"over {lineage.attempts} proposals" + ) + print("\nCommitted lineage:") + for version in lineage.versions: + print( + f" v{version.index} ({version.note}): {version.score:.4g} " + f"[{speedups_of(version.evaluation)}]" + ) + print("\nBest kernel:") + print((source_of(lineage.best.artifact) or repr(lineage.best.artifact)).rstrip()) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--budget", type=int, default=4, help="Variation steps (AVO arm)" + ) + parser.add_argument( + "--step-budget", + type=int, + default=6, + help="Evaluator calls the agent may make within one variation step", + ) + parser.add_argument( + "--evo", + action="store_true", + help="Run the single-turn control instead: the same domain through " + "optimize_anything, matched on evaluator calls", + ) + parser.add_argument( + "--evo-budget", + type=int, + default=0, + help="Iterations for --evo; 0 matches the AVO arm's evaluator calls " + "(--budget x --step-budget)", + ) + parser.add_argument( + "--transfer", + action="store_true", + help="After evolving, adapt the winning kernel to a sibling task in one step " + "-- the paper's MHA -> GQA episode", + ) + parser.add_argument( + "--selection", + choices=["pareto", "best"], + default="pareto", + help="Candidate selection for the --evo control", + ) + parser.add_argument("--seed", type=int, default=0, help="Seed for the control") + args = parser.parse_args() + + if args.evo: + evo = run_evo(args, random.Random(args.seed)) + print("\n" + "=" * 72) + print( + f"mode: EVO control | seed {evo.seed_score:.6g} -> " + f"best {evo.best_score:.6g} ({evo.proposals} proposals, " + f"{evo.evaluations} evaluator calls)" + ) + print("\nBest kernel:") + print((source_of(evo.best.artifact) or repr(evo.best.artifact)).rstrip()) + return + + started = time.monotonic() + lineage, agent = run_avo(args) + report_avo(lineage, agent, time.monotonic() - started) + + if args.transfer: + before, after = transfer(agent, lineage.best.artifact) + print("\n" + "=" * 72) + print( + f"Transfer to {TRANSFER_TASK.name}, one variation step with the agent's " + f"history intact:" + ) + print(f" reference implementation: {before.score:.4g}") + print(f" adapted kernel: {after.score:.4g}") + print(f" [{speedups_of(after)}]") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/ds1000.py b/docs/source/llm_examples/optimization/ds1000.py new file mode 100644 index 000000000..621912edd --- /dev/null +++ b/docs/source/llm_examples/optimization/ds1000.py @@ -0,0 +1,209 @@ +"""Learn scipy from execution feedback: textual-gradient training on DS-1000. + +The training-loop face of `textgrad.py`, ported from the flagship demo of +`strands-labs/ai-functions` (``memory_backprop_scipy``), with the same eight +real DS-1000 scipy training problems and held-out test problem (see +`ds1000_data.py` for provenance). Where `guidelines.py` applies one human +sentence of feedback once, here the feedback is an *oracle's*: every candidate +solution is executed against the benchmark's own tests, and the pass/fail +verdict -- with the error, test input, and expected-vs-actual values on +failure -- is what backpropagates. + +Three steps: + +1. **Direct test** -- solve the held-out problem with empty memory. +2. **Training** -- solve the eight training problems, backpropagate each + problem's verdict, then fold every parameter's merged gradients in a single + accumulation (``scipy_716`` in the training set teaches the ``minimize()`` + contract the held-out problem needs). +3. **Trained test** -- re-solve the held-out problem under the learned + ``coding_patterns`` / ``common_pitfalls``. + +Demonstrates: +- oracle feedback: `ds1000_data.build_feedback` strings from *executing* model + code -- the loop closes programmatically, no human in it +- batch training: eight recorded roots under one optimizer; per-problem + ``backward(optimizer.graph(output_i), feedback_i)`` accumulating on the two + shared boxes; one ``accumulate`` over a synthetic round node gathering the + roots (the flagship's single-consolidate shape) +- detached evaluation: the two test solves pass ``box.value``, so they are + gradient-free by construction -- the held-out problem can never train itself +- the data/skill module split that keeps the held-out test honest: each + problem's ``code_context`` contains its reference solution, and lives in a + module the solver's system prompt only names + +Run with:: + + python -m effectful.handlers.llm.harness \\ + docs/source/llm_examples/optimization/ds1000.py \\ + --model gpt-5-mini --reasoning-effort low --tool-choice none + +Expect a few minutes: ~10 solver calls plus a backward call per training +problem and one accumulation per parameter. Results are stochastic run to run, +and the model matters: the run above flipped the held-out problem FAIL -> PASS +(the direct attempt returned the whole ``OptimizeResult`` where the test wants +``res.x``; training on ``scipy_716`` taught exactly that contract). A weaker +model (gpt-4o-mini) visibly *learns* -- its trained attempts fix the +L-BFGS-B bounds format its direct attempt crashed on -- but tends to keep +failing the held-out problem on some other of its four simultaneous +requirements. +""" + +import argparse +import textwrap + +from docs.source.llm_examples.optimization.ds1000_data import ( + TEST_PROBLEMS, + TRAIN_PROBLEMS, + ExecutionResult, + build_feedback, + execute_and_test, + extract_solution_code, +) +from docs.source.llm_examples.optimization.textgrad import ( + CallNode, + Parameter, + TextGradOptimizer, +) +from effectful.handlers.llm import Skill +from effectful.ops.semantics import handler + +coding_patterns = Parameter( + "No learned patterns yet.", + description=( + "Concise bullet-point list (MAX 15 items) of general, reusable coding " + "patterns and idioms for scipy/data science. Each bullet should be one " + "sentence. Merge similar patterns into a single bullet. Do not include " + "problem-specific details." + ), +) +common_pitfalls = Parameter( + "No known pitfalls yet.", + description=( + "Concise bullet-point list (MAX 15 items) of common, reusable pitfalls " + "and mistakes to avoid. Each bullet should be one sentence. Merge " + "similar pitfalls into a single bullet. Do not include problem-specific " + "details." + ), +) + + +@Skill.define +def solve( + problem: str, library: str, coding_patterns: str, common_pitfalls: str +) -> str: + """Solve the data science problem below by generating Python code. + + Output ONLY the Python code -- no explanations, no markdown fences. The + code will be inserted directly into an execution environment where + {library}, numpy, and the input variables shown in the problem's + block are already defined: use them as they are, never redefine or + re-create them, and assign the answer to exactly the variable the problem + names (the one marked ``# put solution in this variable``). + + + {coding_patterns} + + + + {common_pitfalls} + + + + {problem} + + + Do not use any tools. + """ + + +def attempt(problem: dict, patterns, pitfalls) -> tuple[str, str, ExecutionResult]: + """One solve of ``problem``, scored by the DS-1000 oracle. + + ``patterns`` / ``pitfalls`` are the `Parameter` boxes during training (the + recording handler notes the use and unwraps them) and plain ``.value`` + strings during evaluation (detached: no edge, no gradient). Returns the raw + model output -- the graph key -- alongside the extracted solution and its + verdict. + """ + raw = solve( + problem=problem["prompt"], + library=problem["library"], + coding_patterns=patterns, + common_pitfalls=pitfalls, + ) + solution = extract_solution_code(raw) + return raw, solution, execute_and_test(solution, problem["code_context"]) + + +def show(tag: str, problem: dict, solution: str, result: ExecutionResult) -> None: + print(f"[{tag}] {problem['id']}: {'PASS' if result.passed else 'FAIL'}") + print(textwrap.indent(solution.strip(), " ")) + if not result.passed and result.error: + print(textwrap.indent(f"error: {result.error.strip()}", " ! ")) + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.parse_args() + test = TEST_PROBLEMS[0] + + # Step 1 -- direct test with empty memory, detached (.value): nothing is + # recorded and no gradient can flow from the held-out problem. + print("== Step 1: direct test (empty memory) ==\n") + _, before_solution, before = attempt( + test, coding_patterns.value, common_pitfalls.value + ) + show("direct", test, before_solution, before) + + # Step 2 -- training: eight independent solves recorded as eight roots. + print("== Step 2: training ==\n") + optimizer = TextGradOptimizer() + attempts: list[tuple[dict, str, str, ExecutionResult]] = [] + with handler(optimizer): + for problem in TRAIN_PROBLEMS: + raw, solution, result = attempt(problem, coding_patterns, common_pitfalls) + attempts.append((problem, raw, solution, result)) + + # Backpropagate each problem's oracle verdict; gradients accumulate on the + # two shared boxes across all eight backward passes. + for problem, raw, solution, result in attempts: + print(f"[train] {problem['id']}: {'PASS' if result.passed else 'FAIL'}") + optimizer.backward( + optimizer.graph(raw), build_feedback(problem, solution, result) + ) + print( + f"\naccumulated gradients: {len(coding_patterns.gradients)} on " + f"coding_patterns, {len(common_pitfalls.gradients)} on common_pitfalls" + ) + + # One merged update per parameter: a synthetic round node gathers the + # eight roots so a single accumulate folds each box's gradients at once. + training_round = CallNode( + skill_name="training round", + children=[optimizer.graph(raw) for _, raw, _, _ in attempts], + ) + optimizer.accumulate(training_round) + + print(f"\nlearned coding_patterns:\n{coding_patterns.value}\n") + print(f"learned common_pitfalls:\n{common_pitfalls.value}\n") + + # Step 3 -- re-test with the learned memory, detached again. + print("== Step 3: trained test ==\n") + _, after_solution, after = attempt( + test, coding_patterns.value, common_pitfalls.value + ) + show("trained", test, after_solution, after) + + verdict = { + (False, True): "FAIL -> PASS (memory-driven improvement)", + (True, True): "PASS -> PASS", + (False, False): "FAIL -> FAIL", + (True, False): "PASS -> FAIL", + }[(before.passed, after.passed)] + print(f"{test['id']}: {verdict}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/ds1000_data.py b/docs/source/llm_examples/optimization/ds1000_data.py new file mode 100644 index 000000000..7ec816a52 --- /dev/null +++ b/docs/source/llm_examples/optimization/ds1000_data.py @@ -0,0 +1,237 @@ +"""DS-1000 scipy problems and their execution oracle, for `ds1000.py`. + +Data and no skills, deliberately: the harness splices a skill's defining +module's source into its system prompt, and each problem's ``code_context`` +contains the DS-1000 *reference solution* (inside ``generate_ans``) alongside +the test harness. Keeping problems in a module the solver never shares -- it +is imported by ``ds1000.py``, so the solver's prompt lists only this module's +*name* -- is what makes the held-out evaluation honest (the module-scoping +lesson of ``autoformalization/auditing.py``). + +Problems are the scipy subset curated by `strands-labs/ai-functions` +(Apache-2.0) for its ``memory_backprop_scipy`` demo, taken verbatim from the +DS-1000 benchmark ("DS-1000: A Natural and Reliable Benchmark for Data Science +Code Generation", Lai et al., arXiv:2211.11501; problems derive from +StackOverflow, CC-BY-SA-4.0). Each has an ``id``, ``library``, a ``prompt`` +(the question with a ``BEGIN SOLUTION`` marker), and a ``code_context`` (the +executable test harness that inserts the candidate and asserts correctness). +``scipy_716`` in the training set teaches the ``minimize()`` contract the +held-out test problem needs. + +The oracle below is pure Python -- no model sits in the scoring path. It +``exec``s model-written code by design (the same caveat +``choreographies/multi_agent_choreography.py`` carries): fine against these +fixed benchmark problems, but worth knowing before pointing it elsewhere. +""" + +from typing import Any + +TRAIN_PROBLEMS: list[dict[str, Any]] = [ + { + "id": "scipy_711", + "library": "Scipy", + "prompt": "Problem:\nI have a set of data and I want to compare which line describes it best (polynomials of different orders, exponential or logarithmic).\nI use Python and Numpy and for polynomial fitting there is a function polyfit(). \nHow do I fit y = Alogx + B using polyfit()? The result should be an np.array of [A, B]\nA:\n\nimport numpy as np\nimport scipy\nx = np.array([1, 7, 20, 50, 79])\ny = np.array([10, 19, 30, 35, 51])\n\n\nresult = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\n\n\ndef generate_test_case(test_case_id):\n def define_test_input(test_case_id):\n if test_case_id == 1:\n x = np.array([1, 7, 20, 50, 79])\n y = np.array([10, 19, 30, 35, 51])\n return x, y\n\n def generate_ans(data):\n _a = data\n x, y = _a\n result = np.polyfit(np.log(x), y, 1)\n return result\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n assert np.allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport numpy as np\nimport scipy\nx, y = test_input\n[insert]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_712", + "library": "Scipy", + "prompt": "Problem:\nI have a set of data and I want to compare which line describes it best (polynomials of different orders, exponential or logarithmic).\nI use Python and Numpy and for polynomial fitting there is a function polyfit(). \nHow do I fit y = A + Blogx using polyfit()? The result should be an np.array of [A, B]\nA:\n\nimport numpy as np\nimport scipy\nx = np.array([1, 7, 20, 50, 79])\ny = np.array([10, 19, 30, 35, 51])\n\n\nresult = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\n\n\ndef generate_test_case(test_case_id):\n def define_test_input(test_case_id):\n if test_case_id == 1:\n x = np.array([1, 7, 20, 50, 79])\n y = np.array([10, 19, 30, 35, 51])\n return x, y\n\n def generate_ans(data):\n _a = data\n x, y = _a\n result = np.polyfit(np.log(x), y, 1)[::-1]\n return result\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n assert np.allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport numpy as np\nimport scipy\nx, y = test_input\n[insert]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_713", + "library": "Scipy", + "prompt": "Problem:\nI have a set of data and I want to compare which line describes it best (polynomials of different orders, exponential or logarithmic).\nI use Python and Numpy and for polynomial fitting there is a function polyfit(). But I found no such functions for exponential and logarithmic fitting.\nHow do I fit y = A*exp(Bx) + C ? The result should be an np.array of [A, B, C]. I know that polyfit performs bad for this function, so I would like to use curve_fit to solve the problem, and it should start from initial guess p0.\nA:\n\nimport numpy as np\nimport scipy.optimize\ny = np.array([1, 7, 20, 50, 79])\nx = np.array([10, 19, 30, 35, 51])\np0 = (4, 0.1, 1)\n\nresult = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nimport scipy.optimize\n\n\ndef generate_test_case(test_case_id):\n\n def define_test_input(test_case_id):\n if test_case_id == 1:\n y = np.array([1, 7, 20, 50, 79])\n x = np.array([10, 19, 30, 35, 51])\n p0 = (4, 0.1, 1)\n return x, y, p0\n\n def generate_ans(data):\n _a = data\n x, y, p0 = _a\n result = scipy.optimize.curve_fit(\n lambda t, a, b, c: a * np.exp(b * t) + c, x, y, p0=p0\n )[0]\n return result\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n assert np.allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport numpy as np\nimport scipy.optimize\nx, y, p0 = test_input\n[insert]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_714", + "library": "Scipy", + "prompt": "Problem:\nI can't figure out how to do a Two-sample KS test in Scipy.\nAfter reading the documentation scipy kstest\nI can see how to test where a distribution is identical to standard normal distribution\nfrom scipy.stats import kstest\nimport numpy as np\nx = np.random.normal(0,1,1000)\ntest_stat = kstest(x, 'norm')\n#>>> test_stat\n#(0.021080234718821145, 0.76584491300591395)\nWhich means that at p-value of 0.76 we can not reject the null hypothesis that the two distributions are identical.\nHowever, I want to compare two distributions and see if I can reject the null hypothesis that they are identical, something like:\nfrom scipy.stats import kstest\nimport numpy as np\nx = np.random.normal(0,1,1000)\nz = np.random.normal(1.1,0.9, 1000)\nand test whether x and z are identical\nI tried the naive:\ntest_stat = kstest(x, z)\nand got the following error:\nTypeError: 'numpy.ndarray' object is not callable\nIs there a way to do a two-sample KS test in Python? If so, how should I do it?\nThank You in Advance\nA:\n\nfrom scipy import stats\nimport numpy as np\nnp.random.seed(42)\nx = np.random.normal(0, 1, 1000)\ny = np.random.normal(0, 1, 1000)\n\nstatistic, p_value = ... # put solution in these variables\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nfrom scipy import stats\n\n\ndef generate_test_case(test_case_id):\n\n def define_test_input(test_case_id):\n if test_case_id == 1:\n np.random.seed(42)\n x = np.random.normal(0, 1, 1000)\n y = np.random.normal(0, 1, 1000)\n elif test_case_id == 2:\n np.random.seed(42)\n x = np.random.normal(0, 1, 1000)\n y = np.random.normal(1.1, 0.9, 1000)\n return x, y\n\n def generate_ans(data):\n _a = data\n x, y = _a\n statistic, p_value = stats.ks_2samp(x, y)\n return [statistic, p_value]\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n np.testing.assert_allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nfrom scipy import stats\nimport numpy as np\nnp.random.seed(42)\nx, y = test_input\n[insert]\nresult = [statistic, p_value]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(2):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_715", + "library": "Scipy", + "prompt": "Problem:\nI can't figure out how to do a Two-sample KS test in Scipy.\nAfter reading the documentation scipy kstest\nI can see how to test where a distribution is identical to standard normal distribution\nfrom scipy.stats import kstest\nimport numpy as np\nx = np.random.normal(0,1,1000)\ntest_stat = kstest(x, 'norm')\n#>>> test_stat\n#(0.021080234718821145, 0.76584491300591395)\nWhich means that at p-value of 0.76 we can not reject the null hypothesis that the two distributions are identical.\nHowever, I want to compare two distributions and see if I can reject the null hypothesis that they are identical, something like:\nfrom scipy.stats import kstest\nimport numpy as np\nx = np.random.normal(0,1,1000)\nz = np.random.normal(1.1,0.9, 1000)\nand test whether x and z are identical\nI tried the naive:\ntest_stat = kstest(x, z)\nand got the following error:\nTypeError: 'numpy.ndarray' object is not callable\nIs there a way to do a two-sample KS test in Python, then test whether I can reject the null hypothesis that the two distributions are identical(result=True means able to reject, and the vice versa) based on alpha? If so, how should I do it?\nThank You in Advance\nA:\n\nfrom scipy import stats\nimport numpy as np\nnp.random.seed(42)\nx = np.random.normal(0, 1, 1000)\ny = np.random.normal(0, 1, 1000)\nalpha = 0.01\n\nresult = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nfrom scipy import stats\n\n\ndef generate_test_case(test_case_id):\n\n def define_test_input(test_case_id):\n if test_case_id == 1:\n np.random.seed(42)\n x = np.random.normal(0, 1, 1000)\n y = np.random.normal(0, 1, 1000)\n elif test_case_id == 2:\n np.random.seed(42)\n x = np.random.normal(0, 1, 1000)\n y = np.random.normal(1.1, 0.9, 1000)\n alpha = 0.01\n return x, y, alpha\n\n def generate_ans(data):\n _a = data\n x, y, alpha = _a\n s, p = stats.ks_2samp(x, y)\n result = p <= alpha\n return result\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n np.testing.assert_array_equal(result, ans)\n return 1\n\n\nexec_context = r"""\nfrom scipy import stats\nimport numpy as np\nx, y, alpha = test_input\n[insert]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(2):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_716", # KEY training example — teaches correct minimize() contract + "library": "Scipy", + "prompt": "Problem:\nAccording to the SciPy documentation it is possible to minimize functions with multiple variables, yet it doesn't tell how to optimize on such functions.\nfrom scipy.optimize import minimize\nfrom math import sqrt, sin, pi, cos\ndef f(c):\n return sqrt((sin(pi/2) + sin(0) + sin(c) - 2)**2 + (cos(pi/2) + cos(0) + cos(c) - 1)**2)\nprint minimize(f, 3.14/2 + 3.14/7)\n\nThe above code does try to minimize the function f, but for my task I need to minimize with respect to three variables, starting from `initial_guess`.\nSimply introducing a second argument and adjusting minimize accordingly yields an error (TypeError: f() takes exactly 2 arguments (1 given)).\nHow does minimize work when minimizing with multiple variables.\nI need to minimize f(a,b,c)=((a+b-c)-2)**2 + ((3*a-b-c))**2 + sin(b) + cos(b) + 4.\nResult should be a list=[a,b,c], the parameters of minimized function.\n\nA:\n\nimport scipy.optimize as optimize\nfrom math import sqrt, sin, pi, cos\n\ninitial_guess = [-1, 0, -3]\n\nresult = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nfrom scipy import optimize\n\n\ndef generate_test_case(test_case_id):\n\n def define_test_input(test_case_id):\n if test_case_id == 1:\n a = [-1, 0, -3]\n return a\n\n def generate_ans(data):\n _a = data\n initial_guess = _a\n\n def g(params):\n a, b, c = params\n return (\n ((a + b - c) - 2) ** 2\n + ((3 * a - b - c)) ** 2\n + np.sin(b)\n + np.cos(b)\n + 4\n )\n\n res = optimize.minimize(g, initial_guess)\n result = res.x\n return result\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n def g(params):\n a, b, c = params\n return (\n ((a + b - c) - 2) ** 2 + ((3 * a - b - c)) ** 2 + np.sin(b) + np.cos(b) + 4\n )\n\n assert abs(g(result) - g(ans)) < 1e-2\n return 1\n\n\nexec_context = r"""\nimport scipy.optimize as optimize\nfrom math import sqrt, sin, pi, cos\ninitial_guess = test_input\n[insert]\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_717", + "library": "Scipy", + "prompt": "Problem:\nHow does one convert a list of Z-scores from the Z-distribution (standard normal distribution, Gaussian distribution) to left-tailed p-values? I have yet to find the magical function in Scipy's stats module to do this, but one must be there.\nA:\n\nimport numpy as np\nimport scipy.stats\nz_scores = np.array([-3, -2, 0, 2, 2.5])\n\np_values = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nimport scipy.stats\n\n\ndef generate_test_case(test_case_id):\n\n def define_test_input(test_case_id):\n if test_case_id == 1:\n a = np.array([-3, -2, 0, 2, 2.5])\n return a\n\n def generate_ans(data):\n _a = data\n z_scores = _a\n temp = np.array(z_scores)\n p_values = scipy.stats.norm.cdf(temp)\n return p_values\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n np.testing.assert_allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport numpy as np\nimport scipy.stats\nz_scores = test_input\n[insert]\nresult = p_values\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, + { + "id": "scipy_718", + "library": "Scipy", + "prompt": "Problem:\nHow does one convert a list of Z-scores from the Z-distribution (standard normal distribution, Gaussian distribution) to left-tailed p-values? Original data is sampled from X ~ N(mu, sigma). I have yet to find the magical function in Scipy's stats module to do this, but one must be there.\nA:\n\nimport scipy.stats\nimport numpy as np\nz_scores = [-3, -2, 0, 2, 2.5]\nmu = 3\nsigma = 4\n\np_values = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport scipy\nfrom scipy import sparse\nimport scipy.stats\nimport copy\nimport io\nfrom scipy import integrate\n\n\ndef generate_test_case(test_case_id):\n def define_test_input(test_case_id):\n if test_case_id == 1:\n z_scores = [-3, -2, 0, 2, 2.5]\n mu = 3\n sigma = 4\n return z_scores, mu, sigma\n\n def generate_ans(data):\n _a = data\n z_scores, mu, sigma = _a\n temp = np.array(z_scores)\n p_values = scipy.stats.norm.cdf(temp)\n return p_values\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n np.testing.assert_allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport numpy as np\nimport scipy.stats\nz_scores, mu, sigma = test_input\n[insert]\nresult = p_values\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, +] + +# 1 test problem: scipy_787 +TEST_PROBLEMS = [ + { + "id": "scipy_787", + "library": "Scipy", + "prompt": "Problem:\n\n\nI am having a problem with minimization procedure. Actually, I could not create a correct objective function for my problem.\nProblem definition\n•\tMy function: yn = a_11*x1**2 + a_12*x2**2 + ... + a_m*xn**2,where xn- unknowns, a_m - coefficients. n = 1..N, m = 1..M\n•\tIn my case, N=5 for x1,..,x5 and M=3 for y1, y2, y3.\nI need to find the optimum: x1, x2,...,x5 so that it can satisfy the y\nMy question:\n•\tHow to solve the question using scipy.optimize?\nMy code: (tried in lmfit, but return errors. Therefore I would ask for scipy solution)\nimport numpy as np\nfrom lmfit import Parameters, minimize\ndef func(x,a):\n return np.dot(a, x**2)\ndef residual(pars, a, y):\n vals = pars.valuesdict()\n x = vals['x']\n model = func(x,a)\n return (y - model)**2\ndef main():\n # simple one: a(M,N) = a(3,5)\n a = np.array([ [ 0, 0, 1, 1, 1 ],\n [ 1, 0, 1, 0, 1 ],\n [ 0, 1, 0, 1, 0 ] ])\n # true values of x\n x_true = np.array([10, 13, 5, 8, 40])\n # data without noise\n y = func(x_true,a)\n #************************************\n # Apriori x0\n x0 = np.array([2, 3, 1, 4, 20])\n fit_params = Parameters()\n fit_params.add('x', value=x0)\n out = minimize(residual, fit_params, args=(a, y))\n print out\nif __name__ == '__main__':\nmain()\nResult should be optimal x array. The method I hope to use is L-BFGS-B, with added lower bounds on x.\n\nA:\n\n\n\nimport scipy.optimize\nimport numpy as np\nnp.random.seed(42)\na = np.random.rand(3,5)\nx_true = np.array([10, 13, 5, 8, 40])\ny = a.dot(x_true ** 2)\nx0 = np.array([2, 3, 1, 4, 20])\nx_lower_bounds = x_true / 2\n\nout = ... # put solution in this variable\nBEGIN SOLUTION\n\n", + "code_context": 'import numpy as np\nimport copy\nimport scipy.optimize\n\n\ndef generate_test_case(test_case_id):\n def define_test_input(test_case_id):\n if test_case_id == 1:\n np.random.seed(42)\n a = np.random.rand(3, 5)\n x_true = np.array([10, 13, 5, 8, 40])\n y = a.dot(x_true**2)\n x0 = np.array([2, 3, 1, 4, 20])\n x_bounds = x_true / 2\n return a, x_true, y, x0, x_bounds\n\n def generate_ans(data):\n _a = data\n a, x_true, y, x0, x_lower_bounds = _a\n\n def residual_ans(x, a, y):\n s = ((y - a.dot(x**2)) ** 2).sum()\n return s\n\n bounds = [[x, None] for x in x_lower_bounds]\n out = scipy.optimize.minimize(\n residual_ans, x0=x0, args=(a, y), method="L-BFGS-B", bounds=bounds\n ).x\n return out\n\n test_input = define_test_input(test_case_id)\n expected_result = generate_ans(copy.deepcopy(test_input))\n return test_input, expected_result\n\n\ndef exec_test(result, ans):\n assert np.allclose(result, ans)\n return 1\n\n\nexec_context = r"""\nimport scipy.optimize\nimport numpy as np\na, x_true, y, x0, x_lower_bounds = test_input\n[insert]\nresult = out\n"""\n\n\ndef test_execution(solution: str):\n code = exec_context.replace("[insert]", solution)\n for i in range(1):\n test_input, expected_result = generate_test_case(i + 1)\n test_env = {"test_input": test_input}\n exec(code, test_env)\n assert exec_test(test_env["result"], expected_result)\n', + }, +] + + +# ── Execution oracle (pure Python, no model calls) ─────────────────────────── + + +import dataclasses +import signal +import threading +import traceback +from contextlib import contextmanager + + +@dataclasses.dataclass +class ExecutionResult: + """Outcome of running one candidate solution against a problem's tests.""" + + passed: bool + error: str | None = None + test_input: str | None = None + expected_output: str | None = None + actual_output: str | None = None + + +@contextmanager +def _timeout(seconds: int = 30): + """Raise ``TimeoutError`` if the wrapped block runs longer than ``seconds``. + + SIGALRM only works on the main thread; elsewhere the block runs unguarded. + """ + if threading.current_thread() is not threading.main_thread(): + yield + return + + def handler(signum, frame): + raise TimeoutError(f"Execution timed out after {seconds}s") + + old = signal.signal(signal.SIGALRM, handler) + signal.alarm(seconds) + try: + yield + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old) + + +def _truncated(obj: object, max_len: int = 200) -> str: + s = repr(obj) + return s[:max_len] + "..." if len(s) > max_len else s + + +def extract_solution_code(raw_output: str) -> str: + """Strip markdown fences and DS-1000 solution markers from model output.""" + text = raw_output + if "```python" in text: + start = text.index("```python") + len("```python") + end = text.index("```", start) if "```" in text[start:] else len(text) + text = text[start:end] + elif "```" in text: + start = text.index("```") + 3 + newline = text.find("\n", start) + if newline != -1: + start = newline + 1 + end = text.index("```", start) if "```" in text[start:] else len(text) + text = text[start:end] + for marker in ("BEGIN SOLUTION", "END SOLUTION", "", ""): + text = text.replace(marker, "") + return text + + +def _assertion_detail( + solution_code: str, code_context: str +) -> tuple[str | None, str | None, str | None]: + """Re-run a failed solution to capture (test input, expected, actual).""" + try: + with _timeout(10): + test_env: dict = {} + exec(code_context, test_env) # noqa: S102 + exec_context = test_env.get("exec_context", "") + generate_test_case = test_env.get("generate_test_case") + if not exec_context or not generate_test_case: + return None, None, None + test_input, expected = generate_test_case(1) + run_env: dict = {"test_input": test_input} + exec(exec_context.replace("[insert]", solution_code), run_env) # noqa: S102 + return ( + _truncated(test_input), + _truncated(expected), + _truncated(run_env.get("result")), + ) + except Exception: # noqa: BLE001 -- best-effort detail capture + return None, None, None + + +def execute_and_test(solution_code: str, code_context: str) -> ExecutionResult: + """Run the DS-1000 test harness against a candidate solution. + + The benchmark's own preamble is executed first and *outside* the verdict's + ``except``: a failure there is a broken problem definition, not a wrong + answer, and scoring the two the same way makes an unimportable module in + `code_context` read as the model failing every attempt at that problem -- + with the resulting FAIL then backpropagated as if it were the model's to + learn from. + """ + test_env: dict = {} + try: + exec(code_context, test_env) # noqa: S102 + except Exception as e: + raise RuntimeError(f"benchmark preamble failed to execute: {e}") from e + + try: + with _timeout(30): + test_env["test_execution"](solution_code) + return ExecutionResult(passed=True) + except TimeoutError as e: + return ExecutionResult(passed=False, error=str(e)) + except AssertionError as e: + test_input, expected, actual = _assertion_detail(solution_code, code_context) + return ExecutionResult( + passed=False, + error=f"Test assertion failed: {e}" if str(e) else "Test assertion failed", + test_input=test_input, + expected_output=expected, + actual_output=actual, + ) + except Exception as e: # noqa: BLE001 -- any solution error is a failure, captured as feedback + tb = traceback.format_exception(type(e), e, e.__traceback__) + short = "".join(tb[-3:]) if len(tb) > 3 else "".join(tb) + return ExecutionResult(passed=False, error=f"{type(e).__name__}: {e}\n{short}") + + +def build_feedback(problem: dict, solution: str, result: ExecutionResult) -> str: + """The oracle's verdict as optimizer feedback, one string per problem.""" + if result.passed: + return ( + f"[{problem['library']}] {problem['id']} SOLVED.\n" + f"Working solution:\n{solution}\n" + f"Remember this pattern for similar future problems." + ) + parts = [f"Error: {result.error}"] + if result.test_input: + parts.append(f"Test input: {result.test_input}") + if result.expected_output: + parts.append(f"Expected output: {result.expected_output}") + if result.actual_output: + parts.append(f"Actual output: {result.actual_output}") + return ( + f"[{problem['library']}] {problem['id']} FAILED.\n" + f"Your code:\n{solution}\n" + "\n".join(parts) + ) diff --git a/docs/source/llm_examples/optimization/guidelines.py b/docs/source/llm_examples/optimization/guidelines.py new file mode 100644 index 000000000..35adf975f --- /dev/null +++ b/docs/source/llm_examples/optimization/guidelines.py @@ -0,0 +1,132 @@ +"""Learn writing guidelines from one piece of feedback, by textual gradients. + +Two jokes are written under a shared ``joke_guidelines`` parameter and composed +into an email under a ``formatting_guidelines`` parameter; a single sentence of +feedback on the *email* is then backpropagated by `textgrad.TextGradOptimizer`: +an internal skill splits the feedback between the email's inputs (the two joke +results and the formatting parameter), the jokes' shares are refined onward to +the joke parameter, and accumulation rewrites both parameters in place -- any +later call passing the same boxes writes under the improved guidelines. (The +boxes live in memory here; to persist them across runs, host them as dataclass +fields of a persistent `~effectful.handlers.llm.types.Agent` and the harness's +``SQLitePersister`` checkpoints them with no extra code.) + +Demonstrates: +- `textgrad.Parameter`: a mutable box passed directly as a skill argument; the + recording handler notes the use and splices the wrapped value into the prompt +- one `textgrad.TextGradOptimizer` in both autograd roles: installed as a + handler it is the tape recording plain skill calls (dataflow edges recovered + from call nesting and argument identity); afterwards, *outside* its own + handler scope, ``step`` backpropagates and accumulates, returning the + walked graph and the routed per-node feedback for inspection +- gradients accumulating on one box (``joke_guidelines``) from two uses +""" + +import argparse + +from docs.source.llm_examples.optimization.textgrad import ( + Parameter, + TextGradOptimizer, +) +from effectful.handlers.llm import Skill +from effectful.ops.semantics import handler + +joke_guidelines = Parameter( + "No specific guidelines yet.", + description="Standing guidelines for writing a good joke.", +) +formatting_guidelines = Parameter( + "No specific guidelines yet.", + description="Standing guidelines for the layout and typography of an email.", +) + + +# The skills declare plain ``str`` parameters: the recording handler unwraps a +# `Parameter` box to its value before the prompt is built, so the model sees +# only the text. +@Skill.define +def joke_writer(topic: str, guidelines: str) -> str: + """Write a short joke about the following topic: "{topic}". + + Follow these guidelines: + + {guidelines} + + + Do not use any tools. Return only the joke. + """ + + +@Skill.define +def email_writer(joke_1: str, joke_2: str, guidelines: str) -> str: + """Write a short email to Jane Doe containing the following two jokes, verbatim: + + Joke 1: {joke_1} + + Joke 2: {joke_2} + + Follow these email formatting guidelines: + + {guidelines} + + + Do not use any tools. Return only the email. + """ + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--feedback", + type=str, + default=( + "Both jokes are too wordy: each should be a single-sentence one-liner. " + "The email should include a title for each joke." + ), + help="Natural-language feedback on the final email", + ) + args = parser.parse_args() + + print(f"joke guidelines (before): {joke_guidelines.value}") + print(f"formatting guidelines (before): {formatting_guidelines.value}\n") + + # Forward pass: ordinary calls recorded by the optimizer-as-handler. + # Passing the boxes (not their .value) records the parameter uses; passing + # the jokes' returned strings onward is what wires the joke calls in as + # the email's children. + optimizer = TextGradOptimizer() + with handler(optimizer): + cat_joke = joke_writer(topic="cats", guidelines=joke_guidelines) + prog_joke = joke_writer(topic="programmers", guidelines=joke_guidelines) + email = email_writer( + joke_1=cat_joke, joke_2=prog_joke, guidelines=formatting_guidelines + ) + + print(f"cat joke: {cat_joke}\n") + print(f"prog joke: {prog_joke}\n") + print(f"email:\n{email}\n") + print(f"feedback: {args.feedback}\n") + + # Backward + accumulate, outside the optimizer's own handler scope so its + # internal skill calls do not join the graph they are optimizing. `grads` + # is the backward pass's ephemeral per-node routed feedback, returned for + # display; the persistent gradients live on the Parameter boxes. + graph, grads = optimizer.step(args.feedback) + + def show(node, depth=0): + pad = " " * depth + print(f"{pad}{node.skill_name}: feedback={grads.get(node, [])}") + for name, box in node.parameters: + print(f"{pad} param {name}: gradients={box.gradients}") + for child in node.children: + show(child, depth + 1) + + print("optimizer graph:") + show(graph) + + print(f"\njoke guidelines (after): {joke_guidelines.value}") + print(f"formatting guidelines (after): {formatting_guidelines.value}") + + +if __name__ == "__main__": + main() diff --git a/docs/source/llm_examples/optimization/kernels.py b/docs/source/llm_examples/optimization/kernels.py new file mode 100644 index 000000000..f1c99fb08 --- /dev/null +++ b/docs/source/llm_examples/optimization/kernels.py @@ -0,0 +1,864 @@ +"""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 Skill + +# 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, size: int | None = None) -> 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. + + ``size`` overrides the task's default length, which is what lets one kernel be + scored on a *vector* of configurations (see `evaluate_kernel`). The seed does not + depend on it, so a smaller configuration is a prefix of a larger one and the + configurations differ only in how much work they ask for. + """ + rng = random.Random(zlib.crc32(task.encode())) + return [ + rng.uniform(-1.0, 1.0) + for _ in range(KERNEL_PERF[task][0] if size is None else size) + ] + + +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: + """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.""" + + @Skill.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, size: int | None = None +) -> 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. + """ + default_size, ceiling = KERNEL_PERF[task] + size = default_size if size is None else size + values = perf_input(task, size) + 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_kernel( + kernel: Kernel, + task: KernelTask, + sizes: collections.abc.Sequence[int] = (), +) -> Evaluation: + """Score a kernel: correctness on the small cases, then speed on one or more timed + configurations. + + This is the scoring half of `evaluate_instruction`, factored out because it is the + whole of what an evaluator has to be. `avo.py` optimizes the kernel *directly* and + hands this function to its agent as a tool, so the two examples share one definition + of what a good kernel is -- the correctness gate, the back-to-back baseline timing, + and the shape of the diagnostics -- rather than each writing their own and inviting + the difference to be mistaken for a result. + + ``sizes`` is the vector of timed configurations; empty means the task's single + default size, which is what `evaluate_instruction` uses and what keeps this + script's behaviour unchanged. With several, the score is the geometric mean of the + per-configuration speedups -- geometric because these are ratios, so a candidate + that doubles one configuration and halves another has not broken even -- and each + configuration also becomes its own `Metric`, which is what lets a Pareto search + keep a candidate that wins only at one size. + """ + 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 + + configs = tuple(sizes) or (KERNEL_PERF[task.name][0],) + measured = [measure_speedup(kernel, task.name, size) for size in configs] + speedups = [speedup for speedup, _ in measured] + correct = passed == len(cases) and all(s > 0.0 for s in speedups) + # Geometric, not arithmetic: these are ratios. Guarded by ``correct`` because a + # single zero would take the whole product to zero anyway -- which is the right + # answer, and is what the gate below already says more legibly. + score = statistics.geometric_mean(speedups) if correct else 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 += [timing for _, timing in measured] + 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()) + ) + # One configuration keeps the original wording verbatim: this script's measured + # numbers were produced under it, and the proposer reads this sentence. + if correct and len(configs) == 1: + verdict = ( + f"correct, and {score:.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" + ) + elif correct: + verdict = ( + f"correct, and {score:.2f}x the reference implementation's speed as a " + f"geometric mean over {len(configs)} configuration(s) (" + + ", ".join(f"{s:.2f}x at n={n}" for s, n in zip(speedups, configs)) + + ") -- the score IS that geometric mean, so a correct but ordinary " + "implementation scores about 1.0 and only a faster one improves" + ) + else: + verdict = ( + f"{passed}/{len(cases)} small cases passed; an incorrect kernel " + f"scores zero no matter how fast it is" + ) + diagnostics.append(Diagnostic("verdict", verdict)) + return Evaluation( + score=score, + metrics=[Metric("score", score)] + + [ + Metric(f"speedup@{size}", speedup if correct else 0.0) + for speedup, size in zip(speedups, configs) + ] + + [Metric("cases_passed", float(passed))], + diagnostics=diagnostics, + ) + + +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. They are narrowed back to the two this + script has always had: `evaluate_kernel` also emits a per-configuration metric, + which on one configuration merely restates ``score`` and would double that + dimension's weight in the frontier-frequency sampling. + """ + 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()), + ], + ) + + evaluation = evaluate_kernel(kernel, task) + return Evaluation( + score=evaluation.score, + metrics=[m for m in evaluation.metrics if m.name in ("score", "cases_passed")], + diagnostics=evaluation.diagnostics, + ) + + +class Proposer: + """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.""" + + @Skill.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..ea3a1c46b --- /dev/null +++ b/docs/source/llm_examples/optimization/library.py @@ -0,0 +1,865 @@ +"""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 ``Skill`` + 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) + +A fourth script implements the *successor* claim. "AVO: Agentic Variation Operators" +(arXiv 2603.24517) argues that the single-turn proposer above is itself the bottleneck: +confined to `Reflect`, the model cannot test a candidate, read its traceback, or revise +before committing. AVO replaces the whole variation operator with an autonomous agent +that holds the evaluator as a tool. `evolve_lineage` below is the small amount of +framework that survives that change, and: + + * `avo.py` -- agentic variation on a kernel, against this engine as the control +""" + +import collections.abc +import contextlib +import dataclasses +import inspect +import linecache +import random +import statistics +import typing + +import pydantic.dataclasses + +from effectful.handlers.llm.harness.provision.litellm import LiteLLMConfigurer +from effectful.ops.semantics import handler + +WORKER_MODEL = "openai/gpt-4.1-mini" + + +@contextlib.contextmanager +def worker(model: str): + """Scope a call to the cheap model the artifact is being optimized *for*.""" + with handler(LiteLLMConfigurer(model=model)): + yield + + +# --------------------------------------------------------------------------- +# 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 Skill 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"], + ) + + +# --------------------------------------------------------------------------- +# The agentic variation operator: single-lineage evolution. +# +# "AVO: Agentic Variation Operators for Autonomous Evolutionary Search" (arXiv +# 2603.24517) makes one change to everything above, and it is a change of *scope* +# rather than of machinery. Where `optimize_anything` decomposes variation as +# ``Vary(P) = Generate(Sample(P))`` and confines the model to a single-turn +# ``Generate``, AVO replaces the whole operator with one autonomous agent run, +# ``Vary(P) = Agent(P, K, f)``: the scoring function ``f`` is a tool the agent may +# call, so it can edit, evaluate, read the diagnostics, and revise before committing +# anything. +# +# Almost none of that needs code here, which is the point. The multi-turn loop the +# paper builds is what a `Skill` call already is, and handing over ``f`` is a `Tool` +# in the skill's lexical scope. What *is* left for a framework is small enough to fit +# below: hold the lineage, verify what the agent returns, and commit it if it earns a +# place. `avo.py` supplies the agent and the domain. +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class Version[A]: + """One committed artifact in a lineage: the paper's git commit, in memory. + + Only *committed* versions are recorded. Everything the agent tried and discarded + on the way stays inside its own variation step -- the paper's "500 optimization + directions" behind 40 versions -- which is why a lineage's length is a count of + successes rather than of work done. + """ + + index: int + artifact: A + evaluation: Evaluation + note: str = "" + + @property + def score(self) -> float: + return self.evaluation.score + + +def commits(candidate: Evaluation, incumbent: Evaluation) -> bool: + """The paper's commit rule: "passes correctness checks and matches or improves + the benchmark score relative to the best committed version so far". + + Two clauses, and the first is doing real work. ``score > 0`` is the correctness + gate -- an evaluator in this directory scores a wrong artifact zero however fast + it is -- and without it the "matches" half of the second clause would happily + commit a broken candidate on top of a broken incumbent, forever. + + "Matches" is kept rather than tightened to a strict improvement, because it is + what the paper says and because a lateral move is how a search leaves a plateau. + `evolve_lineage` separately refuses a candidate that *is* the incumbent, which is + the only way matching could be a pure no-op. + + >>> commits(Evaluation(score=1.2), Evaluation(score=1.0)) + True + >>> commits(Evaluation(score=1.0), Evaluation(score=1.0)) + True + >>> commits(Evaluation(score=0.9), Evaluation(score=1.0)) + False + >>> commits(Evaluation(score=0.0), Evaluation(score=0.0)) + False + """ + return candidate.score > 0.0 and candidate.score >= incumbent.score + + +@dataclasses.dataclass +class Lineage[A]: + """The outcome of a single-lineage run: what was committed, and what it cost. + + ``evaluations`` counts only this loop's own verification calls. The agent's calls + to ``f`` during its variation steps are the domain's to count (see `avo.py`), and + they are the larger number by far -- which is precisely the currency on which an + agentic operator is expensive and a single-turn one is cheap, so a comparison that + quotes one arm's proposals against the other's evaluations is not a comparison. + """ + + versions: list[Version[A]] + attempts: int = 0 + rejected: int = 0 + evaluations: int = 0 + + @property + def seed(self) -> Version[A]: + return self.versions[0] + + @property + def best(self) -> Version[A]: + # A commit never lowers the score, so this is the last version; taken as a + # maximum anyway, so the invariant is enforced rather than assumed. + return max(self.versions, key=lambda v: v.score) + + +def evolve_lineage[A]( + *, + vary: collections.abc.Callable[[A, Evaluation], A], + evaluator: collections.abc.Callable[[A], Evaluation], + seed: A, + budget: int, +) -> Lineage[A]: + """Continuous single-lineage evolution: the paper's outer loop, which is all of + the framework that survives once variation becomes an agent. + + Each step hands the incumbent and its evaluation to ``vary`` -- one autonomous + agent run -- and re-scores whatever comes back. The re-scoring is not + redundant. The agent has already called ``f`` itself, probably several times, but + the number that goes in the lineage has to be the framework's own measurement of + the artifact it was actually given: an agent that reports a score it did not + achieve, or that returns a different kernel than the one it tested, is caught here + rather than being taken at its word. It costs one evaluation per step. + + ``vary`` raising is an ordinary outcome, not a crash: the decode-time gate rejects + malformed artifacts after its retries are exhausted, and the step is simply lost. + """ + counters = {"evaluations": 0} + + def evaluate(artifact: A) -> Evaluation: + counters["evaluations"] += 1 + return evaluator(artifact) + + lineage = Lineage(versions=[Version(0, seed, evaluate(seed), note="seed")]) + print(f"[avo] seed scores {lineage.seed.score:.6g}") + + for step in range(1, budget + 1): + incumbent = lineage.best + lineage.attempts += 1 + try: + child = vary(incumbent.artifact, incumbent.evaluation) + except Exception as exc: + lineage.rejected += 1 + reason = " ".join(str(exc).split())[:200] + print( + f" step {step}: no candidate -- the variation step failed at decode " + f"({type(exc).__name__}: {reason})" + ) + continue + + if artifact_key(child) == artifact_key(incumbent.artifact): + lineage.rejected += 1 + print(f" step {step}: rejected -- returned the incumbent unchanged") + continue + + evaluation = evaluate(child) + if commits(evaluation, incumbent.evaluation): + lineage.versions.append( + Version(len(lineage.versions), child, evaluation, note=f"step {step}") + ) + print( + f" step {step}: COMMITTED v{lineage.versions[-1].index} " + f"{incumbent.score:.6g} -> {evaluation.score:.6g}" + ) + else: + lineage.rejected += 1 + print( + f" step {step}: rejected {evaluation.score:.6g} " + f"against the incumbent's {incumbent.score:.6g}" + ) + + lineage.evaluations = counters["evaluations"] + return lineage + + +# --------------------------------------------------------------------------- +# 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..2f2a8c22b --- /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 ``Skill`` 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 Skill + + +@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: + """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.""" + + @Skill.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 Skill's body: it calls this Skill + 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 + """ + + @Skill.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. + """ + + @Skill.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. + """ + + @Skill.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. + """ + + @Skill.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..df9cead9d --- /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 Skill + + +@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 + + +@Skill.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: + """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.""" + + @Skill.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/optimization/textgrad.py b/docs/source/llm_examples/optimization/textgrad.py new file mode 100644 index 000000000..ce1f85419 --- /dev/null +++ b/docs/source/llm_examples/optimization/textgrad.py @@ -0,0 +1,662 @@ +"""Textual gradients: backprop-style credit assignment over live skill calls. + +Implements the core of TextGrad ("Automatic 'Differentiation' via Text", +arXiv:2406.07496), in the form popularized by optimizer/memory libraries such as +`strands-labs/ai-functions`: named text *parameters* feed forward passes, and after a +run an optimizer walks the computation graph backwards, splitting natural-language +feedback across each call's inputs ("textual gradients") and then folding the +gradients accumulated on each parameter into an improved value (`accumulate`). +The mental model is PyTorch autograd -- a `Parameter` is a learnable weight, a +traced run is a forward pass, and `TextGradOptimizer.step` is ``loss.backward()`` +plus ``optimizer.step()``. + +Where `library.py` (optimize_anything) improves an artifact by frontier *search* -- +propose, evaluate, keep the Pareto survivors -- this module improves parameters by +*credit assignment*: one piece of downstream feedback is routed backwards through the +graph to whichever inputs earned it. The two are complementary faces of text +optimization, and each falls out of ordinary effectful idioms rather than a subsystem: + + * **The graph is recorded live by a handler, not reconstructed from an event log.** + Reference implementations rebuild the computation graph post-hoc from a + coordinator's event log, which forces a schema'd memory backend (values must + serialize into recall events and rehydrate later). Here every skill call already + flows through the interceptable operations of + `effectful.handlers.llm.harness.hooks`: the optimizer implements `call_agent` + to record `CallNode` s as calls happen, holding *live references*, and observes + `call_user` / `call_assistant` / `call_tool` (forwarding untouched) to keep + each node's transcript. Nesting gives parent/child edges for free (a skill + called during another's completion loop is its child), and sibling dataflow -- + one call's output passed as another's argument -- is recovered by matching + argument identity against recorded outputs. + + * **A parameter is just a mutable box.** With live references there is nothing to + (de)serialize, so the whole memory-backend apparatus collapses into `Parameter`: + ``value`` (≈ ``.data``, mutated in place by `accumulate`) plus ``gradients`` + (≈ ``.grad``). There is no name and no ``requires_grad``: a use site is named by + the argument the box was passed as, and opting out of optimization is passing + ``param.value`` instead of ``param`` -- ``.value`` is literally ``.detach()``. + Boxes hosted as dataclass fields of a persistent + `~effectful.handlers.llm.types.Agent` are checkpointed by the existing + `~effectful.handlers.llm.harness.durability.persistence.SQLitePersister` + with no code here. + + * **Backward-pass inputs are typed values, not rendered prompts.** The backward + model is itself a `Skill`; its targets arrive as a ``list[Target]`` and the + node's transcript as the captured messages, all through the `Encodable` + bridge -- the "side information is just a typed value" idiom of `library.py`. + + * **Routing is certified at decode time.** The backward skill returns + ``list[Feedback]``, and `Feedback.__post_init__` rejects any ``target`` naming + no input of the node under distribution, so a hallucinated route raises during + decoding and the harness's ``TenacityRetryer`` feeds the error back. Target + names are ephemeral to a single backward call -- a join key, not a graph id. + + * **Dynamic scoping replaces "do not trace yourself".** Record the forward pass + inside ``with handler(optimizer)`` and call ``step`` outside it, and the + optimizer's internal skill calls never enter the graph; no special-case scope + suppression exists. Likewise, *lexical* scoping keeps those internal skills out + of everyone's toolset: they are methods (bound in no module scope, reachable + through no `Agent`), so the harness never advertises them as callable tools. + +The recording assumes a single-threaded forward pass, and results are matched to +arguments by object identity -- interpolating a result into an f-string before +passing it on still computes the right value but drops the dataflow edge. + +Unlike its sibling examples, this module imports and intercepts harness +operations -- its subject *is* the agent loop. Every interception forwards, so +it supplements the launcher's stack rather than shadowing it. + +See `guidelines.py` for the runnable demo. The two are separate modules +deliberately: the harness splices a skill's defining module into its system prompt, +so demo skills sharing this file would read these optimizer prompts and vice versa. +""" + +import contextvars +import dataclasses +import graphlib +import inspect +import typing + +from effectful.handlers.llm import Skill +from effectful.handlers.llm.harness.hooks import ( + Message, + call_agent, + call_assistant, + call_tool, + call_user, +) +from effectful.ops.semantics import fwd +from effectful.ops.syntax import ObjectInterpretation, implements + +# ── Parameters ──────────────────────────────────────────────────────────────── + + +@dataclasses.dataclass(eq=False) +class Parameter[T]: + """A learnable value: the store, the recalled view, and the graph node in one. + + Pass the box itself to a recorded skill call and the optimizer's tracing + handler records the use (under the argument's name) and splices ``value`` + into the prompt; pass ``param.value`` instead to detach -- the value still + flows, but no edge is recorded and no gradient will reach the box. After + ``step``, accumulation rewrites ``value`` in place, so every later use + sees the improved value. (Outside the optimizer's handler scope there is + nothing to unwrap the box, so pass ``param.value`` there too.) + + ``description`` is an optional escape hatch: the backward and accumulation + models already see the traced prompt in which the box's value appeared, so + its role is usually inferable from context; when set, it is passed to both + as explicit format/merge instructions. + + ``eq=False`` keeps identity semantics: two boxes are never "the same + parameter" by value, deduplication is by ``id``, and the box stays hashable + (so it may be a class-level default, shared across instances -- assigning + one in a class body also triggers `__set_name__`, naming the box after the + attribute). + """ + + value: T + description: str = "" + gradients: list[str] = dataclasses.field(default_factory=list) + + def __set_name__(self, owner: type, name: str) -> None: + self.__name__ = name + + def __str__(self) -> str: + """Render the wrapped value (note: f-stringing a box drops its edge).""" + return str(self.value) + + +# ── Graph nodes and argument scanning ──────────────────────────────────────── + + +@dataclasses.dataclass(eq=False) +class CallNode: + """One traced skill call: its transcript, inputs, output, and children. + + Pure record of the forward pass -- the backward pass reads it and writes + nothing here. The only mutable optimization state anywhere is + `Parameter.gradients` / `Parameter.value`; feedback in flight between + nodes lives in a map local to one `TextGradOptimizer.backward` call. + """ + + skill_name: str + messages: list[Message] = dataclasses.field(default_factory=list) + value: typing.Any = None + parameters: list[tuple[str, "Parameter[typing.Any]"]] = dataclasses.field( + default_factory=list + ) + children: list["CallNode"] = dataclasses.field(default_factory=list) + + +def _scan( + value: typing.Any, + label: str, + seen: set[int], + on_parameter: typing.Callable[[str, "Parameter[typing.Any]"], None], + on_value: typing.Callable[[typing.Any], None], +) -> None: + """Walk an argument, reporting `Parameter` boxes and every reachable object. + + Descends through dicts, sequences, sets and dataclass instances (so a box + hosted on an ``Agent`` passed as ``self`` is found under its field name). + ``on_value`` is called with each object so the tracer can match it against + recorded outputs; ``seen`` keeps shared/cyclic structures finite. + """ + if id(value) in seen: + return + seen.add(id(value)) + if isinstance(value, Parameter): + on_parameter(label, value) + return + on_value(value) + if isinstance(value, dict): + for item in value.values(): + _scan(item, label, seen, on_parameter, on_value) + elif isinstance(value, (list, tuple, set, frozenset)): + for item in value: + _scan(item, label, seen, on_parameter, on_value) + elif dataclasses.is_dataclass(value) and not isinstance(value, type): + for f in dataclasses.fields(value): + _scan(getattr(value, f.name), f.name, seen, on_parameter, on_value) + + +def _strip(value: typing.Any) -> typing.Any: + """Replace `Parameter` boxes with their values, rebuilding containers.""" + if isinstance(value, Parameter): + return value.value + if isinstance(value, dict): + return {k: _strip(v) for k, v in value.items()} + if isinstance(value, tuple): + items = (_strip(v) for v in value) + return type(value)(*items) if hasattr(value, "_fields") else tuple(items) + if isinstance(value, (list, set, frozenset)): + return type(value)(_strip(v) for v in value) + return value + + +# ── Model-boundary types of the backward pass ──────────────────────────────── + +_VALID_TARGETS: contextvars.ContextVar[frozenset[str] | None] = contextvars.ContextVar( + "_VALID_TARGETS", default=None +) + + +@dataclasses.dataclass(frozen=True) +class Target: + """One routable input of the node under distribution, as the model sees it. + + Ephemeral to a single backward call: ``name`` is a per-call join key (the + argument name a box was passed as, or a child call's skill name), not a + graph identifier, and nothing here is stored in the graph. + + Deliberately *not* a supertype of `Parameter`: a target is the rendered + snapshot of an input at the moment of one backward call -- and half the + inputs are not parameters at all (a child `CallNode` appears with kind + ``result``). It carries exactly what the live objects don't: a per-call + name and a ``str``-rendered value. In autograd terms, a `Target` is the + argument the backward function is shown; a `Parameter` is the leaf the + routed gradient lands on. + """ + + name: str + kind: typing.Literal["parameter", "result"] + description: str + value: str + + +@dataclasses.dataclass(frozen=True) +class Feedback: + """One routed piece of feedback the backward model emits. + + ``target`` must name a `Target` of the same backward call; construction + checks it against the routing map the optimizer put in scope, so a + hallucinated target raises during decoding and the harness's retryer feeds + the error back to the model. + """ + + target: str + feedback: str + + def __post_init__(self) -> None: + valid = _VALID_TARGETS.get() + if valid is not None and self.target not in valid: + raise ValueError( + f"Feedback references unknown target {self.target!r}. " + f"Valid targets: {sorted(valid)}. Only use names from the " + f"listed inputs." + ) + + +# ── The optimizer ──────────────────────────────────────────────────────────── + + +class TextGradOptimizer(ObjectInterpretation): + """Record a forward pass as a computation graph, then optimize its parameters. + + One object plays both autograd roles: installed as a handler it is the + *tape*, recording every skill call; afterwards ``step`` backpropagates + feedback through the recorded graph and updates the parameters:: + + optimizer = TextGradOptimizer() + with handler(optimizer): # record + draft = write(topic="cats", guidelines=guidelines) # a Parameter + final = polish(draft=draft) + graph = optimizer.step("Too long; cut the preamble.") # optimize + + **Recording.** Edges come from two sources: *nesting* (a skill called while + another's completion loop is open becomes its child) and *sibling dataflow* + (an argument that is -- by object identity -- the output of an earlier + recorded call grafts that call in as a child, unlisting it as a root). + `Parameter` boxes found in the arguments are recorded under the argument + (or dataclass field) name and unwrapped to their values before the call + proceeds. The ``call_user`` / ``call_assistant`` / ``call_tool`` handlers + only observe: they forward untouched and append the produced messages to + the innermost open node's transcript. + + Every call is recorded, learnable or not: a parameterless node may still be + a conduit to boxes in nested calls below it, a detached (``.value``) branch + is still part of what happened, and the tape doubles as an honest trace for + inspection. Which subtrees a gradient can actually reach is decided + retrospectively, by `_toposort`'s pruning at backward time -- recording + never anticipates what the optimizer will find relevant. + + **Optimizing.** ``step`` (and the finer-grained ``backward`` / + ``accumulate`` / ``zero_grad``, which operate on any `CallNode` graph) + invoke the LLM through the skill methods below, so they must run under the + harness but *outside* this object's own handler scope -- otherwise the + optimization would record itself into the graph it is optimizing. The + skills are *methods* rather than module-level functions so they appear in + no other skill's lexical scope, and this class is not an `Agent` so there + is no shared ``__history__``: every backward call is a fresh conversation. + + Recording is single-threaded, and roots accumulate for the lifetime of the + instance -- use one optimizer per forward pass, or pass ``output=`` to + ``step`` to select among several recorded roots. + """ + + def __init__(self) -> None: + self._roots: list[CallNode] = [] + self._stack: list[CallNode] = [] + self._index: dict[int, CallNode] = {} + + # -- Recording (handler methods) ------------------------------------------ + + @implements(call_agent) + def _trace_call(self, skill, *args, **kwargs): + node = CallNode(skill_name=skill.__name__) + + try: + bound = inspect.signature(skill).bind(*args, **kwargs) + bound.apply_defaults() + arguments = dict(bound.arguments) + except TypeError: # let the call itself report the signature error + arguments = {} + + def on_parameter(label: str, box: Parameter[typing.Any]) -> None: + name = label or getattr(box, "__name__", "") or "parameter" + node.parameters.append((name, box)) + + def on_value(value: typing.Any) -> None: + child = self._index.get(id(value)) + if child is None or child is node: + return + if any(c is child for c in node.children): + return + self._roots = [r for r in self._roots if r is not child] + node.children.append(child) + + seen: set[int] = set() + for name, value in arguments.items(): + _scan(value, name, seen, on_parameter, on_value) + + if node.parameters: + args = tuple(_strip(a) for a in args) + kwargs = {k: _strip(v) for k, v in kwargs.items()} + + parent = self._stack[-1] if self._stack else None + self._stack.append(node) + try: + result = fwd(skill, *args, **kwargs) + finally: + self._stack.pop() + + node.value = result + self._index[id(result)] = node + if parent is not None: + parent.children.append(node) + else: + self._roots.append(node) + return result + + @implements(call_user) + def _trace_user(self, *args, **kwargs): + message = fwd(*args, **kwargs) + if self._stack: + self._stack[-1].messages.append(message) + return message + + @implements(call_assistant) + def _trace_assistant(self, *args, **kwargs): + message, tool_calls, result = fwd(*args, **kwargs) + if self._stack: + self._stack[-1].messages.append(message) + return (message, tool_calls, result) + + @implements(call_tool) + def _trace_tool(self, *args, **kwargs): + message, result, is_final = fwd(*args, **kwargs) + if self._stack: + self._stack[-1].messages.append(message) + return (message, result, is_final) + + # -- Backprop ----------------------------------------------------------- + + @staticmethod + def _toposort(root: CallNode) -> list[CallNode]: + """`CallNode` s in reverse topological order (root first), pruned. + + Consumers precede producers, so feedback distributed at a node is + already complete when the walk reaches its children. Subtrees that + reach no `Parameter` are skipped entirely -- the requires-grad + reachability check: no gradient can land there, so no backward call + should be spent there. + + The result answers every graph question the optimizer has. + Distribution order is the list itself; "which children are routable + targets" is membership (a pruned subtree is exactly one that reaches + no parameter, so for any child of a kept node, kept ⟺ routable); and + every reachable `Parameter` lives on a kept node, so `zero_grad` and + `accumulate` need no unpruned walk. + + Collection and pruning are a plain reachability walk (`graphlib` has + no notion of skipping subtrees), but the *ordering* is delegated to + `graphlib.TopologicalSorter` -- a child produces a value its parent + consumes, so children are the parent's dependencies, ``static_order`` + yields producers first, and the backward pass wants the reverse. A + diamond that somehow degenerated into a cycle raises ``CycleError`` + here instead of walking in a silently wrong order. + """ + reaches: dict[int, bool] = {} + + def _reaches_parameter(node: CallNode) -> bool: + """Memoized; seeded ``False`` so a cycle back to ``node`` terminates.""" + nid = id(node) + if nid in reaches: + return reaches[nid] + reaches[nid] = False + reaches[nid] = bool(node.parameters) or any( + _reaches_parameter(c) for c in node.children + ) + return reaches[nid] + + edges: dict[CallNode, list[CallNode]] = {} + stack = [root] + while stack: + node = stack.pop() + if node in edges: + continue + edges[node] = [c for c in node.children if _reaches_parameter(c)] + stack.extend(edges[node]) + return list(graphlib.TopologicalSorter(edges).static_order())[::-1] + + def _targets_of( + self, node: CallNode, kept: set[int] + ) -> tuple[list[Target], dict[str, "Parameter[typing.Any] | CallNode"]]: + """The node's routable inputs, named uniquely for one backward call. + + ``kept`` is the pruned reachable set from `_toposort`: for a child of a + node in it, membership is exactly "leads to a parameter", so children + outside it are not offered as targets (no gradient could land there). + """ + routing: dict[str, Parameter[typing.Any] | CallNode] = {} + targets: list[Target] = [] + + def unique(name: str) -> str: + candidate, n = name, 1 + while candidate in routing: + n += 1 + candidate = f"{name}_{n}" + return candidate + + seen_boxes: set[int] = set() + for label, box in node.parameters: + if id(box) in seen_boxes: + continue + seen_boxes.add(id(box)) + name = unique(label) + routing[name] = box + targets.append( + Target( + name=name, + kind="parameter", + description=box.description, + value=str(box.value), + ) + ) + for child in node.children: + if id(child) not in kept: + continue + name = unique(child.skill_name) + routing[name] = child + targets.append( + Target(name=name, kind="result", description="", value=str(child.value)) + ) + return targets, routing + + @Skill.define + def _compute_gradients( + self, + targets: list[Target], + trace: list[Message], + output: str, + feedback: list[str], + ) -> list[Feedback]: + """You are an optimization agent analyzing one step of an AI workflow to decide + how its inputs should change. You are given the step's routable inputs, the + conversation trace of the step, the step's output, and feedback (issues) that + this output contributed to. Distribute the feedback across the inputs. + + # Inputs + + {targets} + + # Conversation trace + + {trace} + + # Output + + {output} + + # Issues + + {feedback} + + # Rules + + 1. For an input of kind "parameter", the value is a standing instruction that + will be rewritten once using your feedback and then reused on different + future inputs. Phrase feedback as general guidance in this bullet format: + - add: text to add to the value + - update: text to change in the value + - delete: text to remove from the value + 2. For an input of kind "result", the value is the output of an upstream step + that will receive your feedback and re-distribute it to its own inputs. + Phrase feedback as: how this specific result should change to resolve the + issues. + 3. Give feedback to as few inputs as possible: only those whose change would + actually resolve an issue, and only feedback relevant to that input (and to + its description, when present). It is fine to ignore issues that concern no + input. + 4. Feedback for a "parameter" must be general and applicable to future inputs, + never specific to one run's data. + """ + + def backward(self, root: CallNode, feedback: str) -> dict[CallNode, list[str]]: + """Route ``feedback`` from ``root`` through the graph to the parameters. + + The only state this writes is `Parameter.gradients`, which accumulates + deliberately (across ``backward`` calls, and across uses of one box in + several nodes) until `zero_grad`. The per-node feedback in flight is + working state local to this call, returned for inspection and otherwise + discarded -- nothing on the graph itself changes. + + Each node with incoming feedback is distributed exactly once: the walk + is reverse-topological, so by a node's turn every downstream consumer + (all of them, in a diamond) has already routed its refined share here. + """ + order = self._toposort(root) + kept = {id(node) for node in order} + grads: dict[CallNode, list[str]] = {root: [feedback]} + + for node in order: + incoming = grads.get(node) + if not incoming: + continue + targets, routing = self._targets_of(node, kept) + if not routing: + # Nothing routable here: pass the feedback through unchanged. + for child in node.children: + grads.setdefault(child, []).extend(incoming) + continue + + # One backward model call per node, certified against its targets. + token = _VALID_TARGETS.set(frozenset(routing)) + try: + feedbacks = self._compute_gradients( + targets=targets, + trace=list(node.messages), + output=str(node.value), + feedback=list(incoming), + ) + finally: + _VALID_TARGETS.reset(token) + + # Feedback routed to a Parameter lands on the box (the persistent + # ``.grad``); feedback routed to a child call joins the in-flight + # map, re-distributed when the walk reaches that child. + for fb in feedbacks: + target = routing[fb.target] + if isinstance(target, Parameter): + target.gradients.append(fb.feedback) + else: + grads.setdefault(target, []).append(fb.feedback) + return grads + + @Skill.define + def _accumulate[T](self, value: T, gradients: list[str], description: str) -> T: + """Update the value below by applying the accumulated feedback in + ``gradients``; ``description``, when non-empty, says what the value + must contain and how updates should be merged into it -- follow it. + + + {value} + + + + {gradients} + + + + {description} + + + Preserve whatever the feedback does not ask to change, and keep the + format and shape of the current value. Return only the value's new + content -- no commentary about the update itself, no delimiters around + it. Do not use any tools. + """ + + def accumulate(self, root: CallNode) -> None: + """Fold each parameter's accumulated gradients into its value, in place. + + A box recorded by several nodes is updated exactly once, on the union + of its gradients (deduplication is by box identity). The box's fields + travel as the skill's separate arguments, and the reply is decoded + directly as the value's own type: `_accumulate` is generic in the value + (``value: T``, returning ``T``), so the harness infers the + instantiation from the value actually passed and structured decoding + produces the updated value itself. + """ + seen: set[int] = set() + boxes: list[Parameter[typing.Any]] = [] + for node in self._toposort(root): + for _, box in node.parameters: + if id(box) not in seen: + seen.add(id(box)) + boxes.append(box) + for box in boxes: + if not box.gradients: + continue + box.value = self._accumulate(box.value, box.gradients, box.description) + + def zero_grad(self, root: CallNode) -> None: + """Clear the gradients of every `Parameter` reachable from ``root``. + + The pruned walk suffices: a pruned subtree is precisely one containing + no parameters, so every reachable box lives on a kept node. + """ + for node in self._toposort(root): + for _, box in node.parameters: + box.gradients.clear() + + def graph(self, output: typing.Any | None = None) -> CallNode: + """The recorded graph rooted at the call that produced ``output``. + + With no argument, the sole recorded root -- after a linear pipeline, + the final output's node. When several independent roots were recorded + (a training batch, say), ``output`` selects among them by identity. + + This is what makes per-root `backward` composable beyond `step`'s + one-shot form: a batch run records one root per example, calls + ``backward(optimizer.graph(out_i), feedback_i)`` for each, and then + folds every parameter's merged gradients in a single `accumulate`. + """ + if output is not None: + root = self._index.get(id(output)) + if root is None: + raise ValueError( + "output was not produced by a call recorded on this optimizer" + ) + return root + if len(self._roots) == 1: + return self._roots[0] + raise ValueError( + f"optimizer has recorded {len(self._roots)} roots; pass output= " + f"to select the value the feedback is about" + ) + + def step( + self, feedback: str, output: typing.Any | None = None + ) -> tuple[CallNode, dict[CallNode, list[str]]]: + """Backward + accumulate in one call. + + The entry point is `graph`'s resolution of ``output``. Returns that + root together with `backward`'s per-node routed feedback, for + inspection. + """ + root = self.graph(output) + grads = self.backward(root, feedback) + self.accumulate(root) + return root, grads 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..a45d2c8f0 --- /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 skill 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 skill). +""" + +import argparse + +from effectful.handlers.llm import Skill + + +@Skill.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 + """ + + +@Skill.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 + """ + + +@Skill.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 + """ + + +@Skill.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..0586805ce --- /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 Skill + + +@Skill.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/continual.py b/docs/source/llm_examples/reasoning/continual.py new file mode 100644 index 000000000..03b8c979a --- /dev/null +++ b/docs/source/llm_examples/reasoning/continual.py @@ -0,0 +1,215 @@ +"""A self-improving agent: the harness lives on `self`, the transcript is disposable. + +Implements the inference-time loop shared by "Harness RL is Meta-Learning: +Training to Self-Improve at Test Time" (COLM 2026 submission) and "Continual +Harness: Online Adaptation for Self-Improving Foundation Agents" (arXiv +2605.09998) over the unmodified harness stack. Both papers have an agent revise +its own harness -- system prompt, sub-agents, skills, memory -- from experience, +mid-episode, with model weights frozen. Here that harness is the `Player` +agent's instance attributes: + +- ``instructions`` is the mutable half of the system prompt (the papers' `p`), + spliced into every request via the ``{self.instructions}`` hole below; +- ``notes`` is memory (`M`), catalogued by id and title in every request via + ``{self.memory_catalog}`` and read in full through the REPL on demand; +- skills (`K`) and sub-agents (`G`) are whatever `Tool`s and `Skill`s the model + binds onto ``self`` in the REPL -- bound is offered. + +Every revision the papers route through meta-tools is an assignment to ``self`` +in ``exec_code``, and their Refiner is not a component but a *moment*: the +compaction call where the agent promotes what matters onto ``self`` and discards +the transcript it no longer needs, atomically, with the REPL tool's +``compact="conversation"`` mode. That scope, rather than ``compact="turn"``, is the +one this task needs: a `Player`'s history spans every ``plan_presses`` call, so +dropping only the current call's own rounds would free almost nothing. +The schedule for that moment is itself harness content -- a sentence in the +instructions, revisable like everything else. + +The task is a button corridor: each room's door opens after a hidden button +sequence, discoverable by trial (a wrong press resets the room). Codes are +stable within a run, so what the agent writes onto ``self`` -- discovered +codes, a discovery strategy, a solver skill -- genuinely transfers across +compactions: the harness carries the continuity, the transcript is +disposable. + +Conditions (`--condition`): ``scratch`` starts with empty instructions and the +compaction protocol only; ``expert`` starts with a hand-written strategy, +the papers' hand-engineered-harness baseline. The minimalist baseline needs no +code at all: run with ``--tool-choice none`` and the model must answer directly. +""" + +import argparse +import copy +import dataclasses +import random +import typing + +from docs.source.llm_examples.reasoning.gridworlds import Corridor +from effectful.handlers.llm import Agent, Skill + + +@dataclasses.dataclass(frozen=True) +class Note: + """One entry of the agent's memory: catalogued by id and title in every + request, body read on demand through the REPL.""" + + id: str + title: str + body: str + importance: float = 1.0 + + +@dataclasses.dataclass +class Player: + """You are playing a button corridor. Each room's door opens only after its + hidden button sequence is entered; a wrong press resets that room's + progress; the sequences are stable for the whole game. Your score is the + total number of presses, so rediscovering what you already learned is the + main way to lose. Use the REPL to improve yourself as you go: + + - assign to `self.instructions` to rewrite your own standing strategy; + - append `Note`s to `self.notes` for facts worth keeping (discovered codes, + failed hypotheses, where you left off); read one back with + `print(self.notes[i].body)`; + - define reusable functions or `Skill`s in the REPL and assign them onto + `self` (use the function's own name: `self.next_guess = next_guess`) -- + anything bound on `self` is offered to you as a tool on later calls, + whereas one merely defined in the session is gone when you answer; if one + stops earning its keep, `del`ete it. + + Do that writing *before* you answer, in the same call you learned it. The + observation you are shown reports only recent events, so a press whose + result you never wrote down is a press you will pay to make again. If this + request tells you something the last one did not -- a button that clicked, a + button that buzzed, a room completed -- record it, then answer. + + Compact as you go: once the transcript has served its purpose -- say, when a + room's code is recorded in a note -- write what matters onto `self`, then + call `exec_code` with `compact="conversation"`. That leaves the request, and + the call you made it in: your message, the snippet and its output. Every + earlier call goes, and `self` is untouched. So a code recorded in a note is + a code you keep; a code you only ever read off the transcript is a code you + will pay to rediscover. + """ + + instructions: str = "" + notes: list[Note] = dataclasses.field(default_factory=list) + + @property + def memory_catalog(self) -> str: + entries = sorted(self.notes, key=lambda n: -n.importance) + return ( + "\n".join(f"- [{n.id}] {n.title}" for n in entries) + or "(no notes recorded yet)" + ) + + @Skill.define + def plan_presses(self, observation: str) -> list[str]: + """Decide the next button presses to attempt, given the current state + of the corridor: + + + {observation} + + + Your standing strategy (rewrite it via `self.instructions` when you + learn something structural): + + + {self.instructions} + + + Your memory catalog (bodies via `print(self.notes[i].body)` in the REPL): + + + {self.memory_catalog} + + + Return a short list of buttons (each one of "A", "B" or "C") to press + next, in order. Presses are applied until one buzzes, so a plan past + the first uncertain press is wasted only if that press is wrong. + """ + + +def harness_diff(agent: Agent, baseline: dict[str, typing.Any]) -> str: + """A one-line-per-change report of how `agent`'s harness differs from + `baseline` (a ``dict(vars(agent))`` snapshot taken before the run) -- + what the run authored, rebound, and retired, whatever the route.""" + current = {k: v for k, v in vars(agent).items() if not k.startswith("__")} + previous = {k: v for k, v in baseline.items() if not k.startswith("__")} + lines = [] + for name in sorted(current.keys() | previous.keys()): + if name not in previous: + lines.append(f"+ {name} = {current[name]!r}") + elif name not in current: + lines.append(f"- {name}") + elif current[name] is not previous[name] and current[name] != previous[name]: + lines.append(f"~ {name} = {current[name]!r}") + return "\n".join(lines) or "(harness unchanged)" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rooms", type=int, default=6, help="Number of rooms") + parser.add_argument("--length", type=int, default=4, help="Code length per room") + parser.add_argument("--seed", type=int, default=0, help="RNG seed for the codes") + parser.add_argument( + "--budget", type=int, default=400, help="Press budget before giving up" + ) + parser.add_argument( + "--condition", + choices=("scratch", "expert"), + default="scratch", + help="scratch: empty instructions; expert: a hand-written strategy", + ) + args = parser.parse_args() + + # Imported here, not at module scope, for the same reason the expert + # strategy's text lives in `gridworlds` rather than in this file: this + # module's source is embedded in the system prompt, and a module-level + # binding would additionally sit in every Skill's lexical scope -- either + # way the scratch condition would be handed the expert strategy. + from docs.source.llm_examples.reasoning.gridworlds import ( + CORRIDOR_EXPERT_STRATEGY, + ) + + # The answer key. Locals on purpose: module-level names are in every + # Skill's lexical scope, printed into the system prompt, and readable + # through the REPL. + rng = random.Random(args.seed) + codes = [ + "".join(rng.choice(Corridor.BUTTONS) for _ in range(args.length)) + for _ in range(args.rooms) + ] + optimal = sum(len(c) for c in codes) + + player = Player( + instructions=CORRIDOR_EXPERT_STRATEGY if args.condition == "expert" else "" + ) + # Deep, not shallow: `harness_diff` compares the run's attributes against + # this snapshot, and most of what the agent does to its harness it does *in + # place* -- `self.notes.append(...)`. A shallow copy shares that list, so + # the before and after are the same object and every mutation reads as no + # change at all. + baseline = copy.deepcopy(vars(player)) + + corridor = Corridor(codes=codes) + while not corridor.solved and corridor.presses < args.budget: + presses = player.plan_presses(corridor.observe()) + for button in presses: + if corridor.solved or corridor.presses >= args.budget: + break + if button not in corridor.BUTTONS: + corridor.events.append(f"ignored invalid button {button!r}") + continue + corridor.press(button) + + print(f"\nsolved: {corridor.solved}") + print(f"presses: {corridor.presses} (optimal {optimal}, budget {args.budget})") + print(f"\nharness changes this run:\n{harness_diff(player, baseline)}") + + assert corridor.solved, "ran out of press budget before the corridor was solved" + + +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..8dd32af70 --- /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 Skill + + +class CountdownSolver: + """ + 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. + """ + + @Skill.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..9d71162c2 --- /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 Skill + + +@Skill.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..9038a3931 --- /dev/null +++ b/docs/source/llm_examples/reasoning/gridworlds.py @@ -0,0 +1,219 @@ +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 rest of the specification is a comment, not part of the docstring, + # because `__doc__` is *agent-visible*: `world_model_agent.py` passes it to + # the Physicist as the `` it is allowed to know, immediately before + # telling it the dynamics are hidden and must be inferred from the timeline. + # Anything written above this line is handed to the agent; anything below it + # is for the reader. + # + # 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 + + +@dataclasses.dataclass +class Corridor: + """A button corridor: rooms in a row, each behind a hidden button code. + + Not a `Game` -- it observes as prose and acts by button letter rather than + by grid and move -- but the same species of environment: hidden rules, a + visible goal. ``press`` is the whole interface. A correct press advances the + current room's progress ("click"); completing a code opens the door and + moves to the next room; a wrong press resets the room's progress ("buzz"). + + Living in this module rather than the driving script also keeps the + mechanics out of the model's view, as with `PushGame`: the system prompt + embeds the *skill's* module source, not this one's. The codes themselves + are the answer key -- bind them only in locals of the driving script's + ``main()``, never at module scope, which the model can read. + """ + + BUTTONS: typing.ClassVar[str] = "ABC" + + codes: list[str] + room: int = 0 + progress: int = 0 + presses: int = 0 + events: list[str] = dataclasses.field(default_factory=list) + + @property + def solved(self) -> bool: + return self.room >= len(self.codes) + + def press(self, button: str) -> str: + assert not self.solved + self.presses += 1 + if button == self.codes[self.room][self.progress]: + self.progress += 1 + if self.progress == len(self.codes[self.room]): + self.room += 1 + self.progress = 0 + event = ( + f"press {self.presses}: {button} -> CLICK; door {self.room} opens" + ) + else: + event = f"press {self.presses}: {button} -> click (progress {self.progress})" + else: + self.progress = 0 + event = f"press {self.presses}: {button} -> BUZZ; room {self.room} progress reset" + self.events.append(event) + return event + + def observe(self) -> str: + recent = "\n".join(self.events[-12:]) or "(no presses yet)" + return ( + f"room {self.room} of {len(self.codes)} " + f"(codes are {len(self.codes[0])} presses long); " + f"confirmed progress in this room: {self.progress}; " + f"total presses so far: {self.presses}\n" + f"recent events:\n{recent}" + ) + + +# A hand-written reference strategy for `Corridor` -- an expert-harness baseline +# to compare a from-scratch self-improving agent against. It lives here, not in +# the driving script, because the script's own source is embedded in the system +# prompt: a module-level literal there would hand the from-scratch condition the +# expert strategy verbatim and contaminate the comparison. +CORRIDOR_EXPERT_STRATEGY = """\ +Discover each code by prefix extension: with a confirmed prefix P, try P+"A"; +a buzz resets the room, so re-enter P and try P+"B", then P+"C". Every click +extends the confirmed prefix by one. Record each room's confirmed prefix in a +note *immediately* -- after any clear, notes are all you have. Once a room's +full code is known, replay it exactly; never re-derive a recorded code.""" diff --git a/docs/source/llm_examples/reasoning/hanoi.py b/docs/source/llm_examples/reasoning/hanoi.py new file mode 100644 index 000000000..d131bf1ef --- /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 + - Skills defined inside a function that auto-capture closure-scoped tools +""" + +import argparse +import dataclasses +import itertools + +from effectful.handlers.llm import Skill, 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: + + @Skill.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 skill + 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 ``TenacityRetryer``, + 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() + + @Skill.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..cf7ac0403 --- /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 Skill + + +@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 + + +@Skill.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..fe2f90168 --- /dev/null +++ b/docs/source/llm_examples/reasoning/taboo.py @@ -0,0 +1,161 @@ +"""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 Skill, 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 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 + + @Skill.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 that tries to guess the secret word from hints.""" + + @Skill.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", + ) + # Required, with no default: the word is the Guesser's answer key, and a + # default would have to be written here, in the module whose source the + # Guesser reads. `_module_section` puts the *whole* defining module of a + # `Skill` into that agent's system prompt, so a built-in game list would + # hand `make_guess` the very word it is supposed to infer from hints. + parser.add_argument( + "--secret-word", + type=str, + required=True, + metavar="WORD", + help="Secret word the hinter must lead the guesser to without saying it", + ) + parser.add_argument( + "--taboo-words", + nargs="+", + type=str, + required=True, + metavar="WORD", + help="Taboo words the hinter may not say", + ) + args = parser.parse_args() + + print(f"\nGame: '{args.secret_word}' (taboo: {args.taboo_words})") + play_taboo(args.secret_word, args.taboo_words, 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..a74240332 --- /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 Skill + + +@Skill.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..08792b476 --- /dev/null +++ b/docs/source/llm_examples/reasoning/world_model_agent.py @@ -0,0 +1,237 @@ +"""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 ``Skill`` 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 docs.source.llm_examples.reasoning.gridworlds import ( + Action, + Color, + Game, + State, + Transition, +) +from effectful.handlers.llm import Skill + + +@dataclasses.dataclass +class Physicist: + """Reverse-engineers the game by writing its ``step`` rule as Python code.""" + + hint: str + timeline: list[Transition] = dataclasses.field(default_factory=list) + + @Skill.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. + """ + + @Skill.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": + # Imported under the same name as the types at the top of this file. Reaching + # the module a second way would give a second, distinct `State` class, and a + # frozen dataclass only compares equal to its own class -- so `observe()` and + # the synthesized `step` would return incomparable values and every prediction + # check below would report a surprise. + 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..052714ba4 100644 --- a/effectful/handlers/llm/__init__.py +++ b/effectful/handlers/llm/__init__.py @@ -1,3 +1 @@ -from .template import Agent, Template, Tool - -__all__ = ["Agent", "Template", "Tool"] +from .types import * # noqa: F403, F401 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..a2e2ac05a --- /dev/null +++ b/effectful/handlers/llm/harness/__init__.py @@ -0,0 +1,220 @@ +"""Handlers that give the types in :mod:`effectful.handlers.llm.types` their meaning. + +The :func:`harness` function assembles the standard stack; its constituents are +documented in the submodules below and may be recombined or replaced +individually. +""" + +import os +import pathlib +import typing + +import tenacity + +from effectful.handlers.llm.harness.durability.persistence import SQLitePersister +from effectful.handlers.llm.harness.durability.retrying import TenacityRetryer +from effectful.handlers.llm.harness.durability.transaction import HistoryBuilder +from effectful.handlers.llm.harness.execution.builtin import BuiltinExecutor +from effectful.handlers.llm.harness.execution.restricted import ( + RestrictedPythonExecutor, +) +from effectful.handlers.llm.harness.hooks import AgentLoop +from effectful.handlers.llm.harness.legibility.framework import FrameworkDocumenter +from effectful.handlers.llm.harness.legibility.lexical import ( + ImplicitToolExtractor, + LexicalToolExtractor, +) +from effectful.handlers.llm.harness.observability.dump import SystemPromptDumper +from effectful.handlers.llm.harness.observability.langfuse import LangfuseTracer +from effectful.handlers.llm.harness.observability.rich import ( + RichTerminalRenderer, +) +from effectful.handlers.llm.harness.provision.litellm import ( + LiteLLMConfigurer, +) +from effectful.handlers.llm.harness.synthesis.body import ( + FinalBodySynthesizer, +) +from effectful.handlers.llm.harness.synthesis.snippet import StatefulReplSynthesizer +from effectful.handlers.llm.harness.synthesis.toolcall import ( + ExpressionToolCaller, + MixedToolCaller, +) +from effectful.handlers.llm.harness.validation.mypy import MypyTypeChecker +from effectful.handlers.llm.harness.validation.pydantic import PydanticSkillArgValidator +from effectful.handlers.llm.harness.validation.ty import TyTypeChecker +from effectful.ops.semantics import Interpretation, coproduct + + +def harness( + *, + num_retries: int = 5, + langfuse: bool = False, + render: bool = False, + dump_system_prompt: str | os.PathLike[str] | None = None, + persist_db: str | os.PathLike[str] | None = None, + eval_provider: typing.Literal["builtin", "restricted", "none"] = "builtin", + type_checker: typing.Literal["mypy", "ty", "none"] = "ty", + tool_calling: typing.Literal["auto", "code", "json"] = "auto", + tool_collection: typing.Literal["none", "explicit", "auto"] = "explicit", + check_contracts: bool = True, + **provider_config, +) -> Interpretation: + """ + Instantiate the standard `effectful.handlers.llm` handler stack. + Install it with :func:`~effectful.ops.semantics.handler`:: + + with handler(harness(...)): + ... + + 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. `AgentLoop`, the tool pipeline and `LiteLLMConfigurer` -- the agent + loop, the tools it offers from a `Skill`'s lexical scope, and the model + backend it drives. The pipeline is a lexical tool *extractor* chosen by + ``tool_collection`` (`LexicalToolExtractor` for ``"explicit"``, + `ImplicitToolExtractor` for ``"auto"``, nothing for ``"none"``), plus, + above it, the tool *caller* it feeds (`MixedToolCaller` for + ``tool_calling="auto"``, `ExpressionToolCaller` for ``"code"``, none + for ``"json"``). + 2. `FrameworkDocumenter` -- describe the framework's concepts in the system + prompt. + 3. `HistoryBuilder` -- accumulate the message history of a call. + 4. `RichTerminalRenderer` -- live-render the streaming history (if ``render``). + 5. `SystemPromptDumper` -- dump the system prompt (if ``dump_system_prompt``). + 6. The ``type_checker`` and the ``eval_provider`` -- check and run + model-authored Python (each omitted for ``"none"``). + 7. `StatefulReplSynthesizer` and `FinalBodySynthesizer` -- answer a call by + running a snippet, and by synthesizing a function and calling it. Both + are omitted when ``eval_provider="none"``: each advertises a tool + (``exec_code``, ``write_and_run_body``) that only an executor can decode. + 8. `PydanticSkillArgValidator` -- enforce the pre-conditions a caller + wrote into a `Skill`'s parameter annotations (if ``check_contracts``). + 9. `TenacityRetryer` -- retry malformed/failing model output (if + ``num_retries``). + 10. `SQLitePersister` -- checkpoint a persisted `Agent`'s state/history to + SQLite after each successful call (if ``persist_db``). + 11. `LangfuseTracer` -- log calls to Langfuse (if ``langfuse``). + + Args: + num_retries: Attempts for malformed/failing model output (via + `TenacityRetryer`, which is left out of the stack altogether when + this is ``0``) and, independently, for transport-level failures + (via litellm's own ``num_retries``, bound into the request). + 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. + persist_db: If set, path to a SQLite database used to checkpoint a + persisted `~effectful.handlers.llm.types.Agent`'s (one + constructed with an explicit `agent_id`) state and history via + `~effectful.handlers.llm.harness.durability.persistence.SQLitePersister`. + eval_provider: Which provider runs model-authored Python: + ``"builtin"`` (`BuiltinExecutor`, the default), ``"restricted"`` + (`RestrictedPythonExecutor`), or ``"none"`` for no executor -- + which also takes both synthesizers out of the stack, so nothing is + offered that the stack could not then run. + type_checker: Which handler type-checks model-authored Python before it + runs: ``"ty"`` (`TyTypeChecker`, the default), ``"mypy"`` + (`MypyTypeChecker`), or ``"none"`` to run generated code unchecked. + tool_calling: How the model calls the tools in a `Skill`'s lexical + scope. ``"auto"`` (the default) installs `MixedToolCaller`, which + picks per tool: schema-constrained JSON arguments for every tool a + JSON schema can describe faithfully, and the code pathway for the + rest (generic, variadic, or unadvertisable signatures). ``"code"`` + installs `ExpressionToolCaller`: uniformly, the model writes a + Python call expression which is type-checked in the Skill's scope + and evaluated. ``"json"`` is the classic JSON-only pathway with no + caller at all (polymorphic tools degrade to untyped argument + schemas there, and unadvertisable ones are skipped with a + warning). ``"auto"`` and ``"code"`` require an eval provider: + combining either with ``eval_provider="none"`` raises `ValueError` + rather than silently degrading. + tool_collection: Which tools are *collected* from a `Skill`'s lexical + scope, as opposed to how they are called. ``"explicit"`` (the + default) installs `LexicalToolExtractor`: the `Tool`/`Skill` + values in scope, and those held by in-scope `Agent`\\ s. ``"auto"`` + installs `ImplicitToolExtractor` instead, which additionally wraps + ordinary functions and methods that look deliberately published -- + public name, docstring, complete annotations (see + `ImplicitToolExtractor._implicit_tool_candidate`) -- with no + ``Tool.define`` decorator; it makes naming conventions + load-bearing (prefix orchestration helpers with ``_`` to keep them + out of the model's hands). ``"none"`` installs no extractor at + all: the model sees only the tools the harness itself injects + (``exec_code``, ``write_and_run_body``), never the surrounding + scope's. + check_contracts: Install `PydanticSkillArgValidator`, so a `Skill`'s + arguments are validated against the pydantic metadata its parameter + annotations carry. On by default, which makes such an annotation + mean the same thing whether a person or a model supplied the + argument. Turning it off leaves a direct Python call unchecked; a + model-supplied argument is still validated as the tool call is + decoded, and metadata on a *return* annotation is enforced by the + decoder either way. + + Raises: + ValueError: If ``tool_calling`` is ``"auto"`` or ``"code"`` and + ``eval_provider`` is ``"none"``. + """ + h: Interpretation = AgentLoop() + + if tool_calling != "json" and eval_provider == "none": + raise ValueError( + f'tool_calling="{tool_calling}" has the model answer by writing ' + f"Python, so it needs an eval provider to run what it writes. Pass " + f'eval_provider="builtin" or "restricted", or tool_calling="json".' + ) + + if tool_calling == "auto": + h = coproduct(h, MixedToolCaller()) + elif tool_calling == "code": + h = coproduct(h, ExpressionToolCaller()) + json_only = tool_calling == "json" + if tool_collection == "explicit": + h = coproduct(h, LexicalToolExtractor(json_only=json_only)) + elif tool_collection == "auto": + h = coproduct(h, ImplicitToolExtractor(json_only=json_only)) + + h = coproduct(h, LiteLLMConfigurer(num_retries=num_retries, **provider_config)) + h = coproduct(h, FrameworkDocumenter()) + h = coproduct(h, HistoryBuilder()) + + if render: + h = coproduct(h, RichTerminalRenderer()) + + if dump_system_prompt: + h = coproduct( + h, + SystemPromptDumper(path=pathlib.Path(dump_system_prompt)), + ) + + if type_checker == "ty": + h = coproduct(h, TyTypeChecker()) + elif type_checker == "mypy": + h = coproduct(h, MypyTypeChecker()) + + if eval_provider == "restricted": + h = coproduct(h, RestrictedPythonExecutor()) + elif eval_provider == "builtin": + h = coproduct(h, BuiltinExecutor()) + + if eval_provider != "none": + h = coproduct(h, StatefulReplSynthesizer()) + h = coproduct(h, FinalBodySynthesizer()) + + if check_contracts: + h = coproduct(h, PydanticSkillArgValidator()) + + if num_retries > 0: + h = coproduct(h, TenacityRetryer(stop=tenacity.stop_after_attempt(num_retries))) + + if persist_db is not None: + h = coproduct(h, SQLitePersister(pathlib.Path(persist_db))) + + if langfuse: + h = coproduct(h, LangfuseTracer()) + + return h diff --git a/effectful/handlers/llm/harness/__main__.py b/effectful/handlers/llm/harness/__main__.py new file mode 100644 index 000000000..12f53113d --- /dev/null +++ b/effectful/handlers/llm/harness/__main__.py @@ -0,0 +1,271 @@ +""" +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 `Skill`/`Agent` into something runnable. This module +factors that stack into a single object, `harness`, so the scripts themselves +carry none of the boilerplate. + +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 are consumed here; other flags pass through to the script unchanged. +""" + +import argparse +import os +import pdb +import runpy +import sys +import textwrap +import typing + +import litellm + +from effectful.handlers.llm.harness import harness +from effectful.ops.semantics import handler + + +def _reasoning_effort_choices() -> list[str] | None: + """The ``reasoning_effort`` values a provider will actually accept. + + Read from litellm's canonical ``REASONING_EFFORT`` alias so the CLI choices + track litellm across upgrades. That alias -- and not the looser ``Literal`` + on ``litellm.completion``'s own signature, which additionally admits + ``"default"`` -- is the set providers are held to: OpenAI rejects + ``"default"`` outright with ``Unsupported value: 'reasoning_effort' does not + support 'default' with this model``. "Let the model decide" is spelled by + *omitting* the parameter, which is what this flag's ``None`` default does. + + Returns ``None`` (leave the flag unrestricted) if the alias isn't a Literal + we can read, so a shape change in litellm degrades to accepting any string + rather than breaking the launcher. + """ + try: + from litellm.types.llms.openai import REASONING_EFFORT + + literals = [v for v in typing.get_args(REASONING_EFFORT) 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. + + ``allow_abbrev=False`` is what makes the split honest. With argparse's default, + any script flag that is a unique prefix of a harness flag is claimed here + instead of being passed through -- a script's ``--mode`` would be read as this + parser's ``--model``, silently overwriting the model *and* dropping the flag the + script needed. + """ + parser = argparse.ArgumentParser( + prog=f"python -m {__spec__.name}" if __spec__ else None, + description=textwrap.dedent(__doc__), + allow_abbrev=False, + ) + 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=( + "Attempts for malformed/failing LLM output, and for transport-level " + "failures (forwarded to litellm as its own num_retries)" + ), + ) + 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; left unset " + "(the model's own default) when not given", + ) + parser.add_argument( + "--eval-provider", + type=str, + default="builtin", + choices=["builtin", "restricted", "none"], + help="Provider that runs model-authored Python", + ) + parser.add_argument( + "--type-checker", + type=str, + default="ty", + choices=["mypy", "ty", "none"], + help="Handler that type-checks model-authored Python before it runs", + ) + parser.add_argument( + "--tool-calling", + type=str, + default="auto", + choices=["auto", "code", "json"], + help=( + "How the model calls lexical tools: JSON arguments where a schema " + "can describe the tool, code for the rest (auto); by writing a " + "type-checked Python call expression uniformly (code); or JSON " + "arguments only (json)" + ), + ) + parser.add_argument( + "--tool-collection", + type=str, + default="explicit", + choices=["none", "explicit", "auto"], + help=( + "Which tools are collected from a Skill's lexical scope: the " + "declared Tool/Skill values (explicit); those plus qualifying " + "plain functions and methods, no Tool.define decorator needed " + "(auto); or nothing from the scope at all, leaving only the " + "harness's own tools (none)" + ), + ) + parser.add_argument( + "--no-check-contracts", + dest="check_contracts", + action="store_false", + help=( + "Do not validate a Skill's arguments against the pydantic metadata " + "on its parameter annotations" + ), + ) + 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 _needs_responses_api(model: str) -> bool: + """Whether `model` must be addressed through the Responses API to use tools. + + OpenAI rejects function tools alongside reasoning on ``/v1/chat/completions`` + for its GPT-5.4-and-later models (*Function tools with reasoning_effort are + not supported ... use /v1/responses*), and the harness always sends tools. + Asked about a model litellm does not classify -- another provider's, or one + newer than the installed litellm knows -- this answers ``False`` and leaves + the model string alone, so an unfamiliar name degrades to today's behaviour + rather than being rewritten on a guess. + """ + try: + from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config + + return OpenAIGPT5Config.is_model_gpt_5_4_plus_model( + model + ) and not OpenAIGPT5Config.is_model_gpt_5_search_model(model) + except Exception: + return False + + +def _provider_config(ns: argparse.Namespace) -> dict[str, typing.Any]: + """The litellm kwargs `LiteLLMConfigurer` is built from. + + An unset ``--reasoning-effort`` is left out of the request entirely rather + than forwarded as a sentinel. Every value this parameter takes is one some + provider rejects -- OpenAI answers ``reasoning_effort`` of ``"default"`` with + a 400 -- and the only universally safe way to say "whatever the model does by + default" is to say nothing. + + Saying nothing has one consequence worth naming, because it is not obvious + and it is why the model string may be rewritten below. litellm routes a chat + completion to the Responses API only when ``reasoning_effort`` is not None + (``litellm.main.responses_api_bridge_check``), so omitting the parameter also + silently opts a GPT-5.4+ model *out* of the endpoint its tool calls require. + The ``openai/responses/`` prefix is litellm's own way to ask for that + endpoint directly, and it does not depend on the reasoning parameter -- which + lets the effort stay at the model's own default instead of being pinned to a + value nobody chose. An explicit ``--reasoning-effort`` needs none of this: it + triggers the bridge by itself, and is left to do so. + """ + model = ns.model + if ns.reasoning_effort is None and _needs_responses_api(model): + model = f"openai/responses/{model}" + + config: dict[str, typing.Any] = {"model": model, "tool_choice": ns.tool_choice} + if ns.reasoning_effort is not None: + config["reasoning_effort"] = ns.reasoning_effort + return config + + +def main(argv: list[str] | None = None) -> None: + litellm.drop_params = True + 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