Setting up a DecisionAI Optimization Model

Quantagonia’s DecisionAI requires users to define their optimization in a single Python file. From now on, we will refer to this file as the “model file”, or model.py. The model file should contain the following components:

• Input Data: the main class that contains the input data for the optimization problem. Note: It is required to have one main input data class, although this class can contain nested other input data classes.

• Model Variables: a pydantic model that defines the variables of the optimization problem.

• Model: a class that sets up the variables, constraints, and the objective of the optimization problem. Moreover, the class provides helper methods to display the solution or to load the input data from a file. Note: a solution to string method is required to be implemented.

We continue by describing each of these components in more detail.

Input Data

The main input data class should contain the data required to formulate the optimization problem. You can have one main input data class, and this class can contain nested other input data classes for better organization.

An input data class should inherit from the decision_ai.InputData class, which inherits from Pydantic’s BaseModel. This provides automatic validation, serialization, and deserialization capabilities.

Serialization Requirements

Important: Your custom InputData class must support Pydantic’s serialization and deserialization methods:

• model_dump() - Serializes the model to a dictionary/JSON

• model_validate(data) - Creates an instance from dictionary/JSON data

These methods are automatically inherited from Pydantic’s BaseModel and should work out of the box. Pydantic will raise an error if the data is not serializable. DecisionAI uses these methods internally to transmit data between client and server.

Example of serialization:

# Your custom InputData class
input_data = MyInputData(param_1=42, param_2={"key": 100})

# Serialization (automatic)
data_dict = input_data.model_dump()  # {"param_1": 42, "param_2": {"key": 100}}

# Deserialization (automatic)
restored = MyInputData.model_validate(data_dict)  # Creates MyInputData instance

Data Types and Structure

Since decision_ai.InputData inherits from Pydantic’s BaseModel, you can use any data types that Pydantic supports. For nested Pydantic models, use other subclasses of InputData for consistency.

The only restriction is that all data must be serializable - avoid arbitrary types like custom classes without proper serialization support.

For a complete list of supported types and features, refer to the Pydantic documentation.

Nested Models

Important: Any nested Pydantic models should inherit from decision_ai.InputData rather than directly from Pydantic’s BaseModel. This ensures that the agent understands the input data requirements.

# ✅ CORRECT: Nested model inherits from InputData
class MachineConfig(InputData):
    setup_time: float
    processing_time: float

class FactoryInputData(InputData):
    machines: dict[str, MachineConfig]  # Uses InputData subclass

# ❌ AVOID: Mixing InputData with BaseModel
class MachineConfig(BaseModel):  # Don't do this
    setup_time: float

class FactoryInputData(InputData):
    machines: dict[str, MachineConfig]  # Inconsistent configuration

Examples

Basic Example with Primitive Types

In this example, we define a class with various data types:

from decision_ai import InputData


class ModelInputData(InputData):
    param_1: int
    param_2: dict[str, int]
    param_3: list[int]
    param_4: list[dict[str, str]]

Example with Pydantic Features

Since InputData inherits from Pydantic’s BaseModel, you can use all Pydantic features:

from decision_ai import InputData
from pydantic import Field, validator
from typing import Optional, Literal
from datetime import datetime

class AdvancedInputData(InputData):
    # Field with validation and description
    num_workers: int = Field(gt=0, description="Number of workers available")

    # Optional field with default
    max_overtime_hours: Optional[float] = Field(default=8.0, ge=0)

    # Literal type for enumeration
    optimization_goal: Literal["minimize_cost", "maximize_efficiency"] = "minimize_cost"

    # Complex nested structure
    worker_skills: dict[str, list[str]] = Field(default_factory=dict)

    # Datetime fields
    planning_start: datetime
    planning_end: datetime

    # Custom validation
    @validator('planning_end')
    def end_after_start(cls, v, values):
        if 'planning_start' in values and v <= values['planning_start']:
            raise ValueError('planning_end must be after planning_start')
        return v

Traveling Salesman Problem

This is a convenient way of defining an instance of the Traveling Salesman Problem. We store the nodes in a list, the edges in a list of tuples, and the distances in a dictionary of tuples to floats.

from decision_ai import InputData


class TSPInputData(InputData):
    my_tsp_nodes: list[str]
    my_tsp_edges: list[tuple[str, str]]
    my_tsp_costs: dict[str, dict[str, float]]

    def is_valid_edge(self, edge: tuple[str, str]) -> bool:
        return edge in self.my_tsp_edges

Note

You can define helper methods in the same input data class to manipulate the data. In the previous example, we define a helper method to check if an edge is part of the graph.

Nested Input Data

Here’s an example showing nested input data classes for a manufacturing optimization problem. We demonstrate both dictionary-based nesting (machines and products) and direct class nesting (factory containing energy). This shows how to structure complex data hierarchically using different nesting approaches.

Note: All nested classes inherit from InputData (not BaseModel) for consistency.

from decision_ai import InputData


class MachineConfig(InputData):
    setup_time: float
    processing_time: float
    maintenance_interval: int
    max_daily_hours: float = 24.0


class ProductRequirements(InputData):
    min_quantity: int
    max_quantity: int
    quality_threshold: float
    machine_compatibility: list[str]


class EnergyConfig(InputData):
    peak_hours_cost: float
    off_peak_hours_cost: float
    max_peak_consumption: float


class FactoryConfig(InputData):
    max_workers: int
    shifts: list[str]
    energy: EnergyConfig
    operating_hours: dict[str, tuple[int, int]]


class ManufacturingInputData(InputData):
    machines: dict[str, MachineConfig]
    products: dict[str, ProductRequirements]
    factory: FactoryConfig
    planning_horizon: int
    max_total_maintenance_hours: float = 48.0

Model Variables

The model variables object should define the variables of the optimization problem. This class inherits from the decision_ai.PulpVariables, which is used to define the variables of the optimization problem.

Users can define the variables as this class’s attributes. The primitive type of a variable is pulp.LpVariable, but dictionaries, lists, and combinations of those can be used to define the variables. E.g., list[str, dict[str, pulp.LpVariable]] is a valid type for a variable.

Each variable must have a static method init_{variable_name} that returns the variable attribute given the input data. If this method is not defined, a NotImplementedError is raised.

The class inherits the method decision_ai.PulpVariables.get_values() which extracts the values of variables based on the solution of the optimization problem and returns a decision_ai.PulpVariablesWithValues instance. This class is a Pydantic model with the same schema, but replacing pulp.LpVariable with float. You can query the values of the variables from the decision_ai.PulpVariablesWithValues instance in the same way as you would do for the decision_ai.PulpVariables instance.

Example:

import pulp

from decision_ai import PulpVariables


class ModelVariables(PulpVariables):
    var_1: pulp.LpVariable | int = Field(..., description="Variable 1")
    var_2: dict[str, pulp.LpVariable | int] = Field(..., description="Variable 2")

    @staticmethod
    def init_var_1(input_: ModelInputData) -> pulp.LpVariable | int:  # noqa: ARG004
        return pulp.LpVariable("var_1", lowBound=0)

    @staticmethod
    def init_var_2(input_: ModelInputData) -> dict[str, pulp.LpVariable | int]:  # noqa: ARG004
        return {"key": pulp.LpVariable("var_2_key", lowBound=0)}

Base Model

The model class inherits from decision_ai.PulpDecisionAIModel, which is the interface for the optimization problem. This parent method provides abstract methods that dictate the methods or attributes that must be implemented or defined in the model class:

• decision_ai.PulpDecisionAIModel.set_up_objective(): Define the objective of the optimization problem.

• decision_ai.PulpDecisionAIModel.solution_to_str(): Convert the solution to a human-readable string.

• decision_ai.PulpDecisionAIModel.variables_class: The variables class that defines the variables of the optimization problem.

The model class inherits methods to formulate, solve, and other helper methods. For each group of constraints, it is required to use the decision_ai.constraint() decorator.

Let’s go through the individual components one by one.

Set up Objective

The function decision_ai.PulpDecisionAIModel.set_up_objective() defines the objective of the optimization problem.

Example:

def set_up_objective(
    self, input_: ModelInputData, prob: pulp.LpProblem, variables: ModelVariables
) -> pulp.LpProblem:
    prob += pulp.lpSum([input_.param_1 * variables.var_1, input_.param_2["key"] * variables.var_2["key"]])

    return prob

Solution to String

The function decision_ai.PulpDecisionAIModel.solution_to_str() converts the solution to a human-readable string.

Example:

def solution_to_str(self, input_: ModelInputData, solution: Solution) -> str:
    return f"Objective value: {solution.objective}\n\nStatus: {solution.status}"

Constraints

Finally, users need to implement the constraints by using the decision_ai.constraint() decorator. Each family of constraints should be implemented in a separate static method (do not pass self as a parameter), and the decorator should wrap the method. The name of the method will be used as the name of the constraint family in the optimization problem.

Example:

    @constraint
    def constraint_1(input_: ModelInputData, variables: ModelVariables) -> ConstraintGenerator:
        yield variables.var_1 >= input_.param_1

    @constraint
    def constraint_2(input_: ModelInputData, variables: ModelVariables) -> ConstraintGenerator:
        yield variables.var_2["key"] >= input_.param_2["key"]

    # ... other methods

Note

Constraints can be inherited from parent models. This is useful when you want to reuse constraints in multiple models. The parent model must inherit from decision_ai.PulpDecisionAIModel. Only one parent model is allowed.

The solve method

The decision_ai.PulpDecisionAIModel.solve() method solves the optimization problem and returns the best solution found. This method has a default implementation in the decision_ai.PulpDecisionAIModel class. By default, Quantagonia’s cloud HybridSolver is used to solve the optimization problem.

def solve(
    self,
    prob: pulp.LpProblem,
    variables: GenericVariables,
    input_: GenericInput,  # noqa: ARG002
    **kwargs,
) -> Solution[GenericVariables]:
    """Solve the optimization problem using the Quantagonia hybrid solver.
    This method can be overriden by the user to implement a custom solver.

    Args:
        prob: The PuLP problem object.
        variables: The variables of the model.
        input_: The input data of the model.
        **kwargs: Optional solver parameters:
            time_limit: Maximum solving time in seconds
            integrality_tolerance: Tolerance for integer variables
            feasibility_tolerance: Tolerance for constraint violations
            relative_gap: Relative optimality gap
            absolute_gap: Absolute optimality gap
            presolve: Enable/disable presolve (boolean)
            seed: Random seed for reproducibility
            heuristics_only: Use only heuristic methods (boolean)
            objective_limit: Upper bound on objective value

    Returns:
        A solution object with the status, objective, variables, and violated constraints.
    """
    invalid_kwargs = [key for key in kwargs if key not in _VALID_KEYWORD_ARGUMENTS]
    if invalid_kwargs:
        error_msg = f"Invalid keyword arguments: {invalid_kwargs}. Valid arguments are: {_VALID_KEYWORD_ARGUMENTS}."
        raise ValueError(error_msg)

    api_key = self._get_api_key()
    solver = self._create_hybrid_solver(api_key, **kwargs)
    prob.solve(solver=solver)
    return self._create_solution(prob, variables)

However, users can override this method to use a different solver. The arguments of the method should remain the same since they are called internally within a ChatSession. Users are also responsible for returning the solution as an instance of decision_ai.Solution. For an example, please refer to the Simple Example.

Note

If you use pulp as a solver, you can make use of the decision_ai.PulpDecisionAIModel._create_solution() method to create the solution object. This extracts the solution values from the variables object, and builds the solution object.