ai-flow-python/app/service/engine/execution_context.py

from __future__ import annotations

import re
from typing import Any, Optional

_ARRAY_INDEX_RE = re.compile(r"^(.+)\[(\d+)\]$")


class ExecutionContext:
    def __init__(self, inputs: dict[str, Any] | None = None):
        self._inputs: dict[str, Any] = inputs or {}
        self._step_results: dict[str, dict[str, Any]] = {}

    @property
    def inputs(self) -> dict[str, Any]:
        return self._inputs

    def put_step_result(self, step_id: str, outputs: dict[str, Any]):
        self._step_results[step_id] = outputs or {}

    def get_step_result(self, step_id: str) -> dict[str, Any] | None:
        return self._step_results.get(step_id)

    def lookup(self, path: str) -> Any:
        if not path:
            return None

        parts = path.split(".")
        if parts[0] == "inputs":
            current: Any = self._inputs
            idx = 1
        else:
            current = self._step_results.get(parts[0])
            idx = 1

        for i in range(idx, len(parts)):
            if current is None:
                return None

            key = parts[i]
            m = _ARRAY_INDEX_RE.match(key)
            if m:
                array_key = m.group(1)
                array_index = int(m.group(2))
                if isinstance(current, dict):
                    current = current.get(array_key)
                else:
                    return None
                if isinstance(current, list):
                    current = current[array_index] if 0 <= array_index < len(current) else None
                else:
                    return None
            else:
                if isinstance(current, dict):
                    current = current.get(key)
                elif isinstance(current, list):
                    try:
                        index = int(key)
                        current = current[index] if 0 <= index < len(current) else None
                    except (ValueError, IndexError):
                        return None
                else:
                    return None

        return current