"""Constraint builders for OR-Tools CP-SAT solver.
This module contains all constraint-building logic, separating concerns
and making the code more maintainable and testable.
"""
from ortools.sat.python import cp_model
from chuk_mcp_solver.models import (
AllDifferentParams,
CircuitParams,
Constraint,
ConstraintKind,
ConstraintSense,
CumulativeParams,
ElementParams,
ImplicationParams,
LinearConstraintParams,
NoOverlapParams,
ReservoirParams,
TableParams,
)
def build_linear_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a linear constraint.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: LinearConstraintParams = constraint.params # type: ignore[assignment]
# Build linear expression (sum of coefficient * variable)
expr = sum(int(term.coef) * var_map[term.var] for term in params.terms)
# Convert RHS to int (OR-Tools requires integer values)
rhs_int = int(params.rhs)
# Add constraint based on sense
if params.sense == ConstraintSense.LESS_EQUAL:
model.Add(expr <= rhs_int).WithName(constraint.id)
elif params.sense == ConstraintSense.GREATER_EQUAL:
model.Add(expr >= rhs_int).WithName(constraint.id)
else: # EQUAL
model.Add(expr == rhs_int).WithName(constraint.id)
def build_all_different_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build an all-different constraint.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: AllDifferentParams = constraint.params # type: ignore[assignment]
vars_list = [var_map[var_id] for var_id in params.vars]
model.AddAllDifferent(vars_list).WithName(constraint.id)
def build_element_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build an element constraint: target = array[index].
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: ElementParams = constraint.params # type: ignore[assignment]
index_var = var_map[params.index_var]
target_var = var_map[params.target_var]
model.AddElement(index_var, params.array, target_var).WithName(constraint.id)
def build_table_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a table constraint: allowed tuples.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: TableParams = constraint.params # type: ignore[assignment]
vars_list = [var_map[var_id] for var_id in params.vars]
model.AddAllowedAssignments(vars_list, params.allowed_tuples).WithName(constraint.id)
def build_implication_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build an implication constraint: if bool_var then nested_constraint.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
Raises:
ValueError: If the nested constraint is not a linear constraint.
"""
params: ImplicationParams = constraint.params # type: ignore[assignment]
# Get the boolean variable
if_var = var_map[params.if_var]
# For now, only support linear constraints in the "then" part
if params.then.kind != ConstraintKind.LINEAR:
raise ValueError(
f"Implication only supports linear constraints in 'then' part, got: {params.then.kind}"
)
# For nested constraint, params.then.params might be a dict
if isinstance(params.then.params, dict):
linear_params = LinearConstraintParams(**params.then.params)
else:
linear_params = params.then.params # type: ignore[assignment]
expr = sum(int(term.coef) * var_map[term.var] for term in linear_params.terms) # type: ignore[union-attr,misc]
# Convert RHS to int
rhs_int = int(linear_params.rhs) # type: ignore[union-attr]
# Add implication using OnlyEnforceIf
if linear_params.sense == ConstraintSense.LESS_EQUAL: # type: ignore[union-attr]
model.Add(expr <= rhs_int).OnlyEnforceIf(if_var).WithName(constraint.id)
elif linear_params.sense == ConstraintSense.GREATER_EQUAL: # type: ignore[union-attr]
model.Add(expr >= rhs_int).OnlyEnforceIf(if_var).WithName(constraint.id)
else: # EQUAL
model.Add(expr == rhs_int).OnlyEnforceIf(if_var).WithName(constraint.id)
def build_cumulative_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a cumulative constraint for resource scheduling.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: CumulativeParams = constraint.params # type: ignore[assignment]
# Get start variables
start_vars = [var_map[var_id] for var_id in params.start_vars]
# Process durations (can be variables or constants)
if isinstance(params.duration_vars[0], str):
duration_vars = [var_map[var_id] for var_id in params.duration_vars] # type: ignore[misc,index]
else:
duration_vars = params.duration_vars # type: ignore[assignment]
# Process demands (can be variables or constants)
if isinstance(params.demand_vars[0], str):
demand_vars = [var_map[var_id] for var_id in params.demand_vars] # type: ignore[misc,index]
else:
demand_vars = params.demand_vars # type: ignore[assignment]
# Create interval variables for cumulative constraint
intervals = []
for i, (start, duration) in enumerate(zip(start_vars, duration_vars, strict=True)):
interval = model.NewIntervalVar(
start, duration, start + duration, f"{constraint.id}_interval_{i}"
)
intervals.append(interval)
# Add cumulative constraint
model.AddCumulative(intervals, demand_vars, params.capacity).WithName(constraint.id)
def build_circuit_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a circuit constraint for routing problems.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: CircuitParams = constraint.params # type: ignore[assignment]
# Convert arcs from (from, to, var_id) to (from, to, literal)
arcs = []
for from_node, to_node, var_id in params.arcs:
literal = var_map[var_id]
arcs.append((from_node, to_node, literal))
model.AddCircuit(arcs).WithName(constraint.id)
def build_reservoir_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a reservoir constraint for inventory management.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: ReservoirParams = constraint.params # type: ignore[assignment]
# Get time variables
time_vars = [var_map[var_id] for var_id in params.time_vars]
# Add reservoir constraint
model.AddReservoirConstraint(
time_vars,
params.level_changes,
params.min_level,
params.max_level,
).WithName(constraint.id)
def build_no_overlap_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a no-overlap constraint for disjunctive scheduling.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
"""
params: NoOverlapParams = constraint.params # type: ignore[assignment]
# Get start variables
start_vars = [var_map[var_id] for var_id in params.start_vars]
# Process durations (can be variables or constants)
if isinstance(params.duration_vars[0], str):
duration_vars = [var_map[var_id] for var_id in params.duration_vars] # type: ignore[misc,index]
else:
duration_vars = params.duration_vars # type: ignore[assignment]
# Create interval variables
intervals = []
for i, (start, duration) in enumerate(zip(start_vars, duration_vars, strict=True)):
interval = model.NewIntervalVar(
start, duration, start + duration, f"{constraint.id}_interval_{i}"
)
intervals.append(interval)
# Add no-overlap constraint
model.AddNoOverlap(intervals).WithName(constraint.id)
# Constraint builder dispatch map
CONSTRAINT_BUILDERS = {
ConstraintKind.LINEAR: build_linear_constraint,
ConstraintKind.ALL_DIFFERENT: build_all_different_constraint,
ConstraintKind.ELEMENT: build_element_constraint,
ConstraintKind.TABLE: build_table_constraint,
ConstraintKind.IMPLICATION: build_implication_constraint,
ConstraintKind.CUMULATIVE: build_cumulative_constraint,
ConstraintKind.CIRCUIT: build_circuit_constraint,
ConstraintKind.RESERVOIR: build_reservoir_constraint,
ConstraintKind.NO_OVERLAP: build_no_overlap_constraint,
}
def build_constraint(
model: cp_model.CpModel,
constraint: Constraint,
var_map: dict[str, cp_model.IntVar],
) -> None:
"""Build a single constraint using the appropriate builder.
Args:
model: The CP-SAT model.
constraint: The constraint definition.
var_map: Mapping from variable ID to CP-SAT variable.
Raises:
ValueError: If constraint kind is not supported.
"""
builder = CONSTRAINT_BUILDERS.get(constraint.kind)
if builder is None:
raise ValueError(f"Unsupported constraint kind: {constraint.kind}")
builder(model, constraint, var_map)
