Allow .compute to take _options

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

Author: Nijat Khanbabayev

ccflow/flow_model.py

@@ -28,7 +28,7 @@
 from pydantic.errors import PydanticSchemaGenerationError, PydanticUndefinedAnnotation
 
 from .base import BaseModel, ContextBase, ResultBase
-from .callable import CallableModel, Flow, GraphDepList, WrapperModel
+from .callable import CallableModel, Flow, FlowOptions, GraphDepList, WrapperModel
 from .context import FlowContext
 from .exttypes import PyObjectPath
 from .local_persistence import register_ccflow_import_path
@@ -913,20 +913,20 @@ def __init__(self, model: CallableModel):
     def _compute_target(self) -> CallableModel:
         return self._model
 
-    def compute(self, context: Any = _UNSET, /, **kwargs) -> Any:
+    def compute(self, context: Any = _UNSET, /, _options: Optional[FlowOptions] = None, **kwargs) -> Any:
         target = self._compute_target
         generated = _generated_model_instance(target)
         if generated is not None:
             built_context = _build_generated_compute_context(generated, context, kwargs)
-            return _maybe_auto_unwrap_external_result(target, target(built_context))
+            return _maybe_auto_unwrap_external_result(target, target(built_context, _options=_options))
 
         if context is not _UNSET and kwargs:
             raise TypeError("compute() accepts either one context object or contextual keyword arguments, but not both.")
         if context is _UNSET:
             built_context = target.context_type.model_validate(kwargs)
         else:
             built_context = context if isinstance(context, ContextBase) else target.context_type.model_validate(context)
-        return _maybe_auto_unwrap_external_result(target, target(built_context))
+        return _maybe_auto_unwrap_external_result(target, target(built_context, _options=_options))
 
     @property
     def context_inputs(self) -> Dict[str, Any]:

ccflow/tests/test_flow_model.py

@@ -22,7 +22,8 @@
     Lazy,
     ModelRegistry,
 )
-from ccflow.evaluators import GraphEvaluator
+from ccflow.callable import FlowOptions
+from ccflow.evaluators import GraphEvaluator, MemoryCacheEvaluator
 
 
 class SimpleContext(ContextBase):
@@ -996,3 +997,84 @@ class StrListContext(ContextBase):
         @Flow.model(context_type=StrListContext)
         def bad(vals: FromContext[list[int]]) -> int:
             return sum(vals)
+
+
+def test_compute_forwards_options_with_custom_evaluator():
+    calls = {"count": 0}
+
+    @Flow.model
+    def counter(value: FromContext[int]) -> int:
+        calls["count"] += 1
+        return value
+
+    cache = MemoryCacheEvaluator()
+    model = counter()
+
+    result1 = model.flow.compute(value=10, _options=FlowOptions(evaluator=cache, cacheable=True))
+    result2 = model.flow.compute(value=10, _options=FlowOptions(evaluator=cache, cacheable=True))
+
+    assert result1.value == 10
+    assert result2.value == 10
+    assert calls["count"] == 1
+
+
+def test_compute_forwards_options_with_graph_evaluator():
+    @Flow.model
+    def source(value: FromContext[int]) -> int:
+        return value * 10
+
+    @Flow.model
+    def root(x: int, bonus: FromContext[int]) -> int:
+        return x + bonus
+
+    model = root(x=source())
+
+    # GraphEvaluator evaluates in topo order; verify _options flows through
+    # and the graph evaluator is actually used (doesn't raise CycleError, computes correctly)
+    result = model.flow.compute(
+        FlowContext(value=3, bonus=7),
+        _options=FlowOptions(evaluator=GraphEvaluator()),
+    )
+
+    assert result.value == 37
+
+
+def test_compute_forwards_options_through_bound_model():
+    calls = {"count": 0}
+
+    @Flow.model
+    def add(a: int, b: FromContext[int]) -> int:
+        calls["count"] += 1
+        return a + b
+
+    cache = MemoryCacheEvaluator()
+    bound = add(a=10).flow.with_inputs(b=5)
+
+    result1 = bound.flow.compute(_options=FlowOptions(evaluator=cache, cacheable=True))
+    result2 = bound.flow.compute(_options=FlowOptions(evaluator=cache, cacheable=True))
+
+    assert result1.value == 15
+    assert result2.value == 15
+    assert calls["count"] == 1
+
+
+def test_compute_forwards_options_for_plain_callable_model():
+    calls = {"count": 0}
+
+    class Counter(CallableModel):
+        offset: int
+
+        @Flow.call
+        def __call__(self, context: SimpleContext) -> GenericResult[int]:
+            calls["count"] += 1
+            return GenericResult(value=context.value + self.offset)
+
+    cache = MemoryCacheEvaluator()
+    model = Counter(offset=5)
+
+    result1 = model.flow.compute(value=10, _options=FlowOptions(evaluator=cache, cacheable=True))
+    result2 = model.flow.compute(value=10, _options=FlowOptions(evaluator=cache, cacheable=True))
+
+    assert result1.value == 15
+    assert result2.value == 15
+    assert calls["count"] == 1