Constraint satisfaction with Z3

Some problems cannot be expressed as a single equation. A farmer has 20 heads and 56 legs shared between chickens and rabbits. Alice and Bob's ages sum to 40, and in four years Alice will be twice Bob's age. These are constraint satisfaction problems: find integer values that satisfy a set of conditions simultaneously.

Why not just ask the LLM?#

An LLM can solve simple word problems by reasoning through them. But it does arithmetic in its head, without guarantees. As the number of unknowns grows, or when constraints interact in non-obvious ways, errors appear. Ask a small model to juggle five variables and three cross-dependencies, and it will often produce a plausible-looking but wrong answer.

Z3 is an SMT solver. It searches systematically through all possible integer assignments and either returns one that satisfies every constraint, or proves that none exists. The answer is guaranteed correct. The LLM's job is only to read each sentence and write down the constraint it implies. Z3 does the search.

The agent#

Three primitives: declare a variable, add a constraint, solve.

# solver.py
import z3
from opensymbolicai.blueprints import PlanExecute
from opensymbolicai.core import decomposition, primitive


class ConstraintSolver(PlanExecute):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._solver = z3.Solver()
        self._vars: dict[str, z3.ArithRef] = {}

    @primitive(read_only=True)
    def add_var(self, name: str) -> str:
        """Declare an integer variable and return its name.

        Example: c = add_var("chickens")
        """
        self._vars[name] = z3.Int(name)
        return name

    @primitive(read_only=True)
    def add_constraint(self, expr: str) -> str:
        """Add a constraint using declared variable names and return the expression.

        Example: c1 = add_constraint("chickens + rabbits == 20")
        Example: c2 = add_constraint("2*chickens + 4*rabbits == 56")
        """
        self._solver.add(eval(expr, {}, self._vars))
        return expr

    @primitive(read_only=True)
    def solve(self) -> dict[str, int]:
        """Find integer values satisfying all constraints.

        Returns a dict mapping each variable name to its value.
        Raises ValueError if no solution exists.
        """
        if self._solver.check() != z3.sat:
            raise ValueError("No solution exists for the given constraints.")
        model = self._solver.model()
        return {name: model[var].as_long() for name, var in self._vars.items()}

add_var registers a Z3 integer symbol. add_constraint evaluates the expression string against the registered variables and passes it to the solver. solve asks Z3 whether any assignment satisfies all constraints and returns it.

Every primitive returns a value, so every plan statement is an assignment, which PlanExecute requires.

What the plan looks like#

Given the farmer puzzle above, the LLM writes:

c = add_var("chickens")
r = add_var("rabbits")
c1 = add_constraint("chickens + rabbits == 20")
c2 = add_constraint("2*chickens + 4*rabbits == 56")
result = solve()

One variable per unknown. One constraint per sentence. solve() last.

Run it#

main.py creates a fresh agent per problem and runs three:

# main.py (excerpt)
for problem in PROBLEMS:
    agent = ConstraintSolver(llm=llm)   # fresh instance per problem
    result = agent.run(problem.task)

uv run main.py

Output:

[Coin counting]
  result:   {'nickels': 15, 'quarters': 5}
  expected: [5, 15]  [PASS]

[Age puzzle]
  result:   {'alice': 28, 'bob': 12}
  expected: [12, 28]  [PASS]

[Work hours]
  result:   {'alice': 50, 'bob': 35, 'carol': 25}
  expected: [25, 35, 50]  [PASS]

What to notice#

Fresh agent per problem. Z3 solver state lives in the instance (self._solver, self._vars). Reusing the same agent carries over variables and constraints from the previous problem, producing wrong results. A new instance means a clean solver.

Variable names in constraints must match add_var names exactly. add_constraint evaluates the expression string against self._vars. If the plan calls add_var("nickels") but writes add_constraint("nickel + quarters == 20"), the eval raises NameError. The decomposition examples in solver.py show the exact naming pattern the LLM should follow.

Three primitives cover any linear integer problem. Declare variables, express conditions as equality or inequality constraints, call solve. The number of unknowns and constraints scales without changing the interface.