344 lines
12 KiB
Python
344 lines
12 KiB
Python
"""Formula engine — DAG, topological sort, and evaluation.
|
|
|
|
Builds a dependency graph of formula fields, detects circular references,
|
|
and evaluates formulas in topological order.
|
|
|
|
Aggregate context (KTD3):
|
|
- ``SUM({f1})`` — f1 is an aggregate reference (entire column)
|
|
- ``{f1} + 1`` — f1 is a row reference (current record's value)
|
|
- ``{f1} + SUM({f2})`` — mixed: row f1 + column f2 sum
|
|
|
|
The engine distinguishes these by checking if a field reference appears
|
|
as a direct argument to an aggregate function.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
from collections import deque
|
|
|
|
from agentkit.bitable.formula.functions import AGGREGATE_FUNCTIONS, FUNCTION_REGISTRY
|
|
from agentkit.bitable.formula.parser import (
|
|
FormulaDepthExceededError,
|
|
FormulaParseError,
|
|
FormulaSecurityError,
|
|
UnknownFunctionError,
|
|
evaluate_ast,
|
|
parse_formula,
|
|
)
|
|
|
|
|
|
class CircularReferenceError(Exception):
|
|
"""Raised when formula fields form a circular dependency."""
|
|
|
|
|
|
class FormulaEngine:
|
|
"""Formula engine: parse, build DAG, detect cycles, evaluate.
|
|
|
|
Usage::
|
|
|
|
engine = FormulaEngine()
|
|
# Register formula fields
|
|
engine.add_formula(field_id="calc", formula="=SUM({src})", field_refs={"src"})
|
|
# Evaluate for a specific record
|
|
result = engine.evaluate(field_id="calc", field_values={"src": [1, 2, 3]})
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
# field_id → (ast_tree, field_mapping, aggregate_refs, row_refs)
|
|
self._formulas: dict[str, _FormulaEntry] = {}
|
|
# DAG: field_id → set of field_ids it depends on
|
|
self._dag: dict[str, set[str]] = {}
|
|
|
|
def add_formula(self, field_id: str, formula: str) -> None:
|
|
"""Register a formula for a field.
|
|
|
|
Raises:
|
|
FormulaParseError: Syntax error.
|
|
FormulaSecurityError: Disallowed AST node.
|
|
UnknownFunctionError: Unregistered function.
|
|
CircularReferenceError: Adding this formula creates a cycle.
|
|
"""
|
|
tree, field_mapping, cross_table_mapping = parse_formula(
|
|
formula, set(FUNCTION_REGISTRY.keys())
|
|
)
|
|
|
|
# Classify field refs into aggregate vs row context
|
|
aggregate_refs, row_refs = _classify_refs(tree, field_mapping)
|
|
|
|
entry = _FormulaEntry(
|
|
tree=tree,
|
|
field_mapping=field_mapping,
|
|
aggregate_refs=aggregate_refs,
|
|
row_refs=row_refs,
|
|
cross_table_mapping=cross_table_mapping,
|
|
formula=formula,
|
|
)
|
|
self._formulas[field_id] = entry
|
|
# Update DAG: this field depends on all referenced fields
|
|
# Cross-table refs depend on the relation field (local) + target field (foreign)
|
|
cross_table_local_deps = {rel_id for rel_id, _ in cross_table_mapping.values()}
|
|
self._dag[field_id] = aggregate_refs | row_refs | cross_table_local_deps
|
|
|
|
# Check for cycles
|
|
cycle = _detect_cycle(self._dag)
|
|
if cycle:
|
|
# Rollback
|
|
del self._formulas[field_id]
|
|
del self._dag[field_id]
|
|
raise CircularReferenceError(f"Circular reference detected: {' → '.join(cycle)}")
|
|
|
|
def remove_formula(self, field_id: str) -> None:
|
|
"""Remove a formula from the engine."""
|
|
self._formulas.pop(field_id, None)
|
|
self._dag.pop(field_id, None)
|
|
# Remove edges pointing to this field
|
|
for deps in self._dag.values():
|
|
deps.discard(field_id)
|
|
|
|
def get_dependencies(self, field_id: str) -> set[str]:
|
|
"""Get the set of field IDs that ``field_id`` depends on."""
|
|
return self._dag.get(field_id, set()).copy()
|
|
|
|
def get_cross_table_mapping(self, field_id: str) -> dict[str, tuple[str, str]] | None:
|
|
"""Get the cross-table reference mapping for a formula field.
|
|
|
|
Returns ``{safe_name: (relation_field_id, target_field_id)}`` or
|
|
``None`` if the field is not registered or has no cross-table refs.
|
|
"""
|
|
entry = self._formulas.get(field_id)
|
|
if entry is None:
|
|
return None
|
|
return entry.cross_table_mapping or None
|
|
|
|
def get_dependents(self, field_id: str) -> set[str]:
|
|
"""Get the set of formula field IDs that depend on ``field_id``."""
|
|
return {fid for fid, deps in self._dag.items() if field_id in deps}
|
|
|
|
def topological_order(self) -> list[str]:
|
|
"""Return all formula field IDs in topological order (Kahn's algorithm)."""
|
|
return _topological_sort(self._dag)
|
|
|
|
def evaluate(
|
|
self,
|
|
field_id: str,
|
|
row_values: dict[str, object],
|
|
column_values: dict[str, list[object]] | None = None,
|
|
cross_table_values: dict[str, list[object]] | None = None,
|
|
) -> object:
|
|
"""Evaluate a formula field for a specific record.
|
|
|
|
Args:
|
|
field_id: The formula field to evaluate.
|
|
row_values: Field ID → value for the current record (row context).
|
|
column_values: Field ID → list of all values in that column
|
|
(aggregate context). Required for aggregate references.
|
|
cross_table_values: Safe-name → list of values from related
|
|
records (cross-table context, U3). The service layer resolves
|
|
these before calling evaluate.
|
|
|
|
Returns:
|
|
The computed value.
|
|
|
|
Raises:
|
|
KeyError: Field ID not registered.
|
|
FormulaParseError: Field reference not found in values.
|
|
"""
|
|
if field_id not in self._formulas:
|
|
raise KeyError(f"Formula not registered: {field_id}")
|
|
|
|
entry = self._formulas[field_id]
|
|
column_values = column_values or {}
|
|
cross_table_values = cross_table_values or {}
|
|
|
|
# Build the field_values dict for the evaluator
|
|
# Aggregate refs get column values (lists), row refs get row values (scalars)
|
|
eval_values: dict[str, object] = {}
|
|
|
|
# Map real field IDs to safe names
|
|
for safe_name, real_id in entry.field_mapping.items():
|
|
if real_id in entry.aggregate_refs:
|
|
eval_values[safe_name] = column_values.get(real_id, [])
|
|
else:
|
|
eval_values[safe_name] = row_values.get(real_id)
|
|
|
|
# Cross-table refs: the service layer pre-resolves these to value lists
|
|
for safe_name in entry.cross_table_mapping:
|
|
if safe_name in cross_table_values:
|
|
eval_values[safe_name] = cross_table_values[safe_name]
|
|
else:
|
|
eval_values[safe_name] = []
|
|
|
|
return evaluate_ast(entry.tree, eval_values, FUNCTION_REGISTRY)
|
|
|
|
def evaluate_all_for_record(
|
|
self,
|
|
row_values: dict[str, object],
|
|
column_values: dict[str, list[object]] | None = None,
|
|
) -> dict[str, object]:
|
|
"""Evaluate all registered formulas for a record.
|
|
|
|
Returns a dict of field_id → computed value.
|
|
Formulas are evaluated in topological order so that formula-to-formula
|
|
dependencies are resolved correctly.
|
|
|
|
List return values (from FILTER/SPLIT/LOOKUP) are stored as-is.
|
|
Downstream formulas referencing a list-valued field receive the list.
|
|
"""
|
|
results: dict[str, object] = {}
|
|
column_values = column_values or {}
|
|
|
|
for field_id in self.topological_order():
|
|
# Include already-computed formula results in row_values
|
|
merged_row = {**row_values, **results}
|
|
try:
|
|
results[field_id] = self.evaluate(field_id, merged_row, column_values)
|
|
except (
|
|
FormulaParseError,
|
|
FormulaSecurityError,
|
|
UnknownFunctionError,
|
|
FormulaDepthExceededError,
|
|
ZeroDivisionError,
|
|
TypeError,
|
|
ValueError,
|
|
) as e:
|
|
results[field_id] = {"__error": str(e)}
|
|
|
|
return results
|
|
|
|
|
|
# ── Internal data structures ──────────────────────────────
|
|
|
|
|
|
class _FormulaEntry:
|
|
"""Parsed formula metadata."""
|
|
|
|
__slots__ = (
|
|
"tree",
|
|
"field_mapping",
|
|
"aggregate_refs",
|
|
"row_refs",
|
|
"cross_table_mapping",
|
|
"formula",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
tree: ast.Expression,
|
|
field_mapping: dict[str, str],
|
|
aggregate_refs: set[str],
|
|
row_refs: set[str],
|
|
cross_table_mapping: dict[str, tuple[str, str]],
|
|
formula: str,
|
|
) -> None:
|
|
self.tree = tree
|
|
self.field_mapping = field_mapping
|
|
self.aggregate_refs = aggregate_refs
|
|
self.row_refs = row_refs
|
|
self.cross_table_mapping = cross_table_mapping
|
|
self.formula = formula
|
|
|
|
|
|
# ── DAG utilities ─────────────────────────────────────────
|
|
|
|
|
|
def _detect_cycle(dag: dict[str, set[str]]) -> list[str] | None:
|
|
"""Detect a cycle in the DAG using DFS. Returns the cycle path or None."""
|
|
WHITE, GRAY, BLACK = 0, 1, 2
|
|
color: dict[str, int] = {node: WHITE for node in dag}
|
|
parent: dict[str, str | None] = {node: None for node in dag}
|
|
|
|
def _dfs(node: str) -> list[str] | None:
|
|
color[node] = GRAY
|
|
for neighbor in dag.get(node, set()):
|
|
if neighbor not in color:
|
|
color[neighbor] = WHITE
|
|
parent[neighbor] = None
|
|
if color[neighbor] == GRAY:
|
|
# Found cycle — reconstruct path
|
|
cycle = [neighbor]
|
|
current = node
|
|
while current is not None and current != neighbor:
|
|
cycle.append(current)
|
|
current = parent.get(current)
|
|
cycle.append(neighbor)
|
|
cycle.reverse()
|
|
return cycle
|
|
if color[neighbor] == WHITE:
|
|
parent[neighbor] = node
|
|
result = _dfs(neighbor)
|
|
if result is not None:
|
|
return result
|
|
color[node] = BLACK
|
|
return None
|
|
|
|
for node in dag:
|
|
if color.get(node, WHITE) == WHITE:
|
|
result = _dfs(node)
|
|
if result is not None:
|
|
return result
|
|
return None
|
|
|
|
|
|
def _topological_sort(dag: dict[str, set[str]]) -> list[str]:
|
|
"""Kahn's algorithm for topological sort."""
|
|
# Build in-degree map
|
|
in_degree: dict[str, int] = {node: 0 for node in dag}
|
|
for node, deps in dag.items():
|
|
for dep in deps:
|
|
if dep in in_degree:
|
|
in_degree[node] += 1
|
|
|
|
# Start with nodes that have no dependencies
|
|
queue = deque(node for node, degree in in_degree.items() if degree == 0)
|
|
result: list[str] = []
|
|
|
|
while queue:
|
|
node = queue.popleft()
|
|
result.append(node)
|
|
# Find all nodes that depend on this node
|
|
for other_node, deps in dag.items():
|
|
if node in deps:
|
|
in_degree[other_node] -= 1
|
|
if in_degree[other_node] == 0 and other_node not in result:
|
|
queue.append(other_node)
|
|
|
|
return result
|
|
|
|
|
|
def _classify_refs(
|
|
tree: ast.Expression, field_mapping: dict[str, str]
|
|
) -> tuple[set[str], set[str]]:
|
|
"""Classify field references into aggregate (column) and row context.
|
|
|
|
A field reference is aggregate if it appears as a direct argument to
|
|
an aggregate function (SUM/AVG/COUNT/MIN/MAX). Otherwise it's row context.
|
|
"""
|
|
aggregate_refs: set[str] = set()
|
|
row_refs: set[str] = set()
|
|
|
|
# Get all safe names → real field IDs
|
|
safe_to_real = field_mapping
|
|
|
|
class _Classifier(ast.NodeVisitor):
|
|
def visit_Call(self, node: ast.Call) -> None:
|
|
if isinstance(node.func, ast.Name) and node.func.id in AGGREGATE_FUNCTIONS:
|
|
# Direct arguments to aggregate functions are column refs
|
|
for arg in node.args:
|
|
if isinstance(arg, ast.Name) and arg.id in safe_to_real:
|
|
aggregate_refs.add(safe_to_real[arg.id])
|
|
else:
|
|
# Non-field args (e.g., literals) — visit normally
|
|
self.visit(arg)
|
|
else:
|
|
self.generic_visit(node)
|
|
|
|
def visit_Name(self, node: ast.Name) -> None:
|
|
if node.id in safe_to_real:
|
|
real_id = safe_to_real[node.id]
|
|
if real_id not in aggregate_refs:
|
|
row_refs.add(real_id)
|
|
|
|
_Classifier().visit(tree)
|
|
return aggregate_refs, row_refs
|