palimpzest 0.7.21__py3-none-any.whl → 0.8.1__py3-none-any.whl

palimpzest/query/optimizer/rules.py

@@ -1,38 +1,72 @@
 import logging
+import os
 from copy import deepcopy
 from itertools import combinations
 
-from palimpzest.constants import AggFunc, Cardinality, PromptStrategy
+from palimpzest.constants import AggFunc, Model, PromptStrategy
+from palimpzest.core.data.context_manager import ContextManager
+from palimpzest.core.lib.schemas import AudioBase64, AudioFilepath, ImageBase64, ImageFilepath, ImageURL
+from palimpzest.prompts import CONTEXT_SEARCH_PROMPT
 from palimpzest.query.operators.aggregate import ApplyGroupByOp, AverageAggregateOp, CountAggregateOp
-from palimpzest.query.operators.code_synthesis_convert import CodeSynthesisConvertSingle
+from palimpzest.query.operators.compute import SmolAgentsCompute
 from palimpzest.query.operators.convert import LLMConvertBonded, NonLLMConvert
 from palimpzest.query.operators.critique_and_refine_convert import CriticAndRefineConvert
+from palimpzest.query.operators.distinct import DistinctOp
 from palimpzest.query.operators.filter import LLMFilter, NonLLMFilter
+from palimpzest.query.operators.join import NestedLoopsJoin
 from palimpzest.query.operators.limit import LimitScanOp
 from palimpzest.query.operators.logical import (
     Aggregate,
     BaseScan,
-    CacheScan,
+    ComputeOperator,
+    ContextScan,
     ConvertScan,
+    Distinct,
     FilteredScan,
     GroupByAggregate,
+    JoinOp,
     LimitScan,
-    MapScan,
     Project,
     RetrieveScan,
+    SearchOperator,
 )
-from palimpzest.query.operators.map import MapOp
 from palimpzest.query.operators.mixture_of_agents_convert import MixtureOfAgentsConvert
+from palimpzest.query.operators.physical import PhysicalOperator
 from palimpzest.query.operators.project import ProjectOp
 from palimpzest.query.operators.rag_convert import RAGConvert
 from palimpzest.query.operators.retrieve import RetrieveOp
-from palimpzest.query.operators.scan import CacheScanDataOp, MarshalAndScanDataOp
+from palimpzest.query.operators.scan import ContextScanOp, MarshalAndScanDataOp
+from palimpzest.query.operators.search import (
+    SmolAgentsSearch,  # SmolAgentsCustomManagedSearch,  # SmolAgentsManagedSearch
+)
 from palimpzest.query.operators.split_convert import SplitConvert
 from palimpzest.query.optimizer.primitives import Expression, Group, LogicalExpression, PhysicalExpression
-from palimpzest.utils.model_helpers import get_models, get_vision_models
 
 logger = logging.getLogger(__name__)
 
+# DEFINITIONS
+IMAGE_LIST_FIELD_TYPES = [
+    list[ImageBase64],
+    list[ImageFilepath],
+    list[ImageURL],
+    list[ImageBase64] | None,
+    list[ImageFilepath] | None,
+    list[ImageURL] | None,
+]
+IMAGE_FIELD_TYPES = IMAGE_LIST_FIELD_TYPES + [
+    ImageBase64, ImageFilepath, ImageURL,
+    ImageBase64 | None, ImageFilepath | None, ImageURL | None,
+]
+AUDIO_LIST_FIELD_TYPES = [
+    list[AudioBase64],
+    list[AudioFilepath],
+    list[AudioBase64] | None,
+    list[AudioFilepath] | None,
+]
+AUDIO_FIELD_TYPES = AUDIO_LIST_FIELD_TYPES + [
+    AudioBase64, AudioFilepath,
+    AudioBase64 | None, AudioFilepath | None,
+]
 
 class Rule:
     """
@@ -43,12 +77,12 @@ class Rule:
     def get_rule_id(cls):
         return cls.__name__
 
-    @staticmethod
-    def matches_pattern(logical_expression: LogicalExpression) -> bool:
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         raise NotImplementedError("Calling this method from an abstract base class.")
 
-    @staticmethod
-    def substitute(logical_expression: LogicalExpression, **kwargs) -> set[Expression]:
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **kwargs: dict) -> set[Expression]:
         raise NotImplementedError("Calling this method from an abstract base class.")
 
 
@@ -59,9 +93,9 @@ class TransformationRule(Rule):
     which are created during the substitution.
     """
 
-    @staticmethod
+    @classmethod
     def substitute(
-        logical_expression: LogicalExpression, groups: dict[int, Group], expressions: dict[int, Expression], **kwargs
+        cls, logical_expression: LogicalExpression, groups: dict[int, Group], expressions: dict[int, Expression], **kwargs
     ) -> tuple[set[LogicalExpression], set[Group]]:
         """
         This function applies the transformation rule to the logical expression, which
@@ -81,15 +115,15 @@ class PushDownFilter(TransformationRule):
     most expensive operator in the input group.
     """
 
-    @staticmethod
-    def matches_pattern(logical_expression: Expression) -> bool:
+    @classmethod
+    def matches_pattern(cls, logical_expression: Expression) -> bool:
         is_match = isinstance(logical_expression.operator, FilteredScan)
         logger.debug(f"PushDownFilter matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
-    @staticmethod
+    @classmethod
     def substitute(
-        logical_expression: LogicalExpression, groups: dict[int, Group], expressions: dict[int, Expression], **kwargs
+        cls, logical_expression: LogicalExpression, groups: dict[int, Group], expressions: dict[int, Expression], **kwargs: dict
     ) -> tuple[set[LogicalExpression], set[Group]]:
         logger.debug(f"Substituting PushDownFilter for {logical_expression}")
 
@@ -113,7 +147,7 @@ class PushDownFilter(TransformationRule):
             #       we see a regression / bug in the future
             for expr in input_group.logical_expressions:
                 # if the expression operator is not a convert or a filter, we cannot swap
-                if not (isinstance(expr.operator, (ConvertScan, FilteredScan))):
+                if not (isinstance(expr.operator, (ConvertScan, FilteredScan, JoinOp))):
                     continue
 
                 # if this filter depends on a field generated by the expression we're trying to swap with, we can't swap
@@ -141,8 +175,8 @@ class PushDownFilter(TransformationRule):
                 group_id, group = None, None
 
                 # if the expression already exists, lookup the group_id and group
-                if new_filter_expr.get_expr_id() in expressions:
-                    group_id = expressions[new_filter_expr.get_expr_id()].group_id
+                if new_filter_expr.expr_id in expressions:
+                    group_id = expressions[new_filter_expr.expr_id].group_id
                     new_filter_expr.set_group_id(group_id)
                     group = groups[group_id]
 
@@ -190,8 +224,7 @@ class PushDownFilter(TransformationRule):
                 # create final new logical expression with expr's operator pulled up
                 new_expr = LogicalExpression(
                     expr.operator.copy(),
-                    input_group_ids=[group_id]
-                    + [g_id for g_id in logical_expression.input_group_ids if g_id != input_group_id],
+                    input_group_ids=[group_id] + [g_id for g_id in logical_expression.input_group_ids if g_id != input_group_id],
                     input_fields=group.fields,
                     depends_on_field_names=expr.depends_on_field_names,
                     generated_fields=expr.generated_fields,
@@ -211,218 +244,280 @@ class ImplementationRule(Rule):
     Base class for implementation rules which convert a logical expression to a physical expression.
     """
 
-    pass
-
-
-class NonLLMConvertRule(ImplementationRule):
-    """
-    Substitute a logical expression for a UDF ConvertScan with a NonLLMConvert physical implementation.
-    """
+    @classmethod
+    def _get_image_fields(cls, logical_expression: LogicalExpression) -> set[str]:
+        """Returns the set of fields which have an image (or list[image]) type."""
+        return set([
+            field_name.split(".")[-1]
+            for field_name, field in logical_expression.input_fields.items()
+            if field.annotation in IMAGE_FIELD_TYPES and field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
     @classmethod
-    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
-        is_match = isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is not None
-        logger.debug(f"NonLLMConvertRule matches_pattern: {is_match} for {logical_expression}")
-        return is_match
+    def _get_list_image_fields(cls, logical_expression: LogicalExpression) -> set[str]:
+        """Returns the set of fields which have a list[image] type."""
+        return set([
+            field_name.split(".")[-1]
+            for field_name, field in logical_expression.input_fields.items()
+            if field.annotation in IMAGE_LIST_FIELD_TYPES and field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
-        logger.debug(f"Substituting NonLLMConvertRule for {logical_expression}")
+    def _get_audio_fields(cls, logical_expression: LogicalExpression) -> set[str]:
+        """Returns the set of fields which have an audio (or list[audio]) type."""
+        return set([
+            field_name.split(".")[-1]
+            for field_name, field in logical_expression.input_fields.items()
+            if field.annotation in AUDIO_FIELD_TYPES and field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-        logical_op = logical_expression.operator
+    @classmethod
+    def _get_list_audio_fields(cls, logical_expression: LogicalExpression) -> set[str]:
+        """Returns the set of fields which have a list[audio] type."""
+        return set([
+            field_name.split(".")[-1]
+            for field_name, field in logical_expression.input_fields.items()
+            if field.annotation in AUDIO_LIST_FIELD_TYPES and field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-        # get initial set of parameters for physical op
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
-            {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
-            }
-        )
+    @classmethod
+    def _is_image_only_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes only image input(s) and False otherwise."""
+        return all([
+            field.annotation in IMAGE_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-        # construct multi-expression
-        op = NonLLMConvert(**op_kwargs)
-        expression = PhysicalExpression(
-            operator=op,
-            input_group_ids=logical_expression.input_group_ids,
-            input_fields=logical_expression.input_fields,
-            depends_on_field_names=logical_expression.depends_on_field_names,
-            generated_fields=logical_expression.generated_fields,
-            group_id=logical_expression.group_id,
-        )
+    @classmethod
+    def _is_image_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes image input(s) and False otherwise."""
+        return any([
+            field.annotation in IMAGE_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-        deduped_physical_expressions = set([expression])
-        logger.debug(f"Done substituting NonLLMConvertRule for {logical_expression}")
+    @classmethod
+    def _is_audio_only_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes only audio input(s) and False otherwise."""
+        return all([
+            field.annotation in AUDIO_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-        return deduped_physical_expressions
+    @classmethod
+    def _is_audio_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes audio input(s) and False otherwise."""
+        return any([
+            field.annotation in AUDIO_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
+    @classmethod
+    def _is_text_only_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes only text input(s) and False otherwise."""
+        return all([
+            field.annotation not in IMAGE_FIELD_TYPES + AUDIO_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
 
-class LLMConvertBondedRule(ImplementationRule):
-    """
-    Substitute a logical expression for a ConvertScan with a bonded convert physical implementation.
-    """
+    @classmethod
+    def _is_text_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes text input(s) and False otherwise."""
+        return any([
+            field.annotation not in IMAGE_FIELD_TYPES + AUDIO_FIELD_TYPES
+            for field_name, field in logical_expression.input_fields.items()
+            if field_name.split(".")[-1] in logical_expression.depends_on_field_names
+        ])
+
+    # TODO: support powerset of text + image + audio (+ video) multi-modal operations
+    @classmethod
+    def _is_text_image_multimodal_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes text and image inputs and False otherwise."""
+        return cls._is_image_operation(logical_expression) and cls._is_text_operation(logical_expression)
 
     @classmethod
-    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
-        is_match = isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is None
-        logger.debug(f"LLMConvertBondedRule matches_pattern: {is_match} for {logical_expression}")
-        return is_match
+    def _is_text_audio_multimodal_operation(cls, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the logical_expression processes text and audio inputs and False otherwise."""
+        return cls._is_audio_operation(logical_expression) and cls._is_text_operation(logical_expression)
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
-        logger.debug(f"Substituting LLMConvertBondedRule for {logical_expression}")
+    def _model_matches_input(cls, model: Model, logical_expression: LogicalExpression) -> bool:
+        """Returns True if the model is capable of processing the input and False otherwise."""
+        # compute how many image fields are in the input, and whether any fields are list[image] fields
+        num_image_fields = len(cls._get_image_fields(logical_expression))
+        has_list_image_field = len(cls._get_list_image_fields(logical_expression)) > 0
+        num_audio_fields = len(cls._get_audio_fields(logical_expression))
+        has_list_audio_field = len(cls._get_list_audio_fields(logical_expression)) > 0
+
+        # corner-case: for now, all operators use text or vision models for processing inputs to __call__
+        if model.is_embedding_model():
+            return False
 
+        # corner-case: Llama vision models cannot handle multiple image inputs (at least using Together)
+        if model.is_llama_model() and model.is_vision_model() and (num_image_fields > 1 or has_list_image_field):
+            return False
+
+        # corner-case: Gemini models cannot handle multiple audio inputs
+        if model.is_vertex_model() and model.is_audio_model() and (num_audio_fields > 1 or has_list_audio_field):
+            return False
+
+        # text-only input and text supporting model
+        if cls._is_text_only_operation(logical_expression) and model.is_text_model():
+            return True
+
+        # image-only input and image supporting model
+        if cls._is_image_only_operation(logical_expression) and model.is_vision_model():
+            return True
+
+        # audio-only input and audio supporting model
+        if cls._is_audio_only_operation(logical_expression) and model.is_audio_model():
+            return True
+
+        # multi-modal input and multi-modal supporting model
+        if cls._is_text_image_multimodal_operation(logical_expression) and model.is_text_image_multimodal_model():  # noqa: SIM103
+            return True
+
+        # multi-modal input and multi-modal supporting model
+        if cls._is_text_audio_multimodal_operation(logical_expression) and model.is_text_audio_multimodal_model():  # noqa: SIM103
+            return True
+
+        return False
+
+    @classmethod
+    def _get_fixed_op_kwargs(cls, logical_expression: LogicalExpression, runtime_kwargs: dict) -> dict:
+        """Get the fixed set of physical op kwargs provided by the logical expression and the runtime keyword arguments."""
+        # get logical operator 
         logical_op = logical_expression.operator
 
-        # get initial set of parameters for physical op
+        # set initial set of parameters for physical op
         op_kwargs = logical_op.get_logical_op_params()
         op_kwargs.update(
             {
-                "verbose": physical_op_params["verbose"],
+                "verbose": runtime_kwargs["verbose"],
                 "logical_op_id": logical_op.get_logical_op_id(),
+                "unique_logical_op_id": logical_op.get_unique_logical_op_id(),
                 "logical_op_name": logical_op.logical_op_name(),
+                "api_base": runtime_kwargs["api_base"],
             }
         )
 
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-        pure_text_models = {model for model in text_models if model not in vision_models}
-        pure_vision_models = {model for model in vision_models if model not in text_models}
-
-        # compute attributes about this convert operation
-        is_image_conversion = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        num_image_fields = sum(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        list_image_field = any(
-            [
-                field.is_image_field and hasattr(field, "element_type")
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
+        return op_kwargs
+
+    @classmethod
+    def _perform_substitution(
+        cls,
+        logical_expression: LogicalExpression,
+        physical_op_class: type[PhysicalOperator],
+        runtime_kwargs: dict,
+        variable_op_kwargs: list[dict] | dict | None = None,
+    ) -> set[PhysicalExpression]:
+        """
+        This performs basic substitution logic which proceeds in four steps:
+
+            1. The basic kwargs for the physical operator are computed using the logical operator
+               and runtime kwargs.
+            2. If variable kwargs are provided, then they are merged with the basic kwargs and one
+               instance of the physical operator is created for each dictionary of variable kwargs.
+            3. A physical expression is created for each physical operator instance.
+            4. The unique set of physical expressions is returned.
+
+        Args:
+            logical_expression (LogicalExpression): The logical expression containing a logical operator.
+            physical_op_class (type[PhysicalOperator]): The class of the physical operator we wish to construct.
+            runtime_kwargs (dict): Keyword arguments which are provided at runtime.
+            variable_op_kwargs (list[dict] | dict | None): A (list of) variable kwargs to customize each
+                physical operator instance.
+
+        Returns:
+            set[PhysicalExpression]: The unique set of physical expressions produced by initializing the
+                physical_op_class with the provided keyword arguments.
+        """
+        # get physical operator kwargs which are fixed for each instance of the physical operator
+        fixed_op_kwargs = cls._get_fixed_op_kwargs(logical_expression, runtime_kwargs)
+
+        # make variable_op_kwargs a list of dictionaries
+        if variable_op_kwargs is None:
+            variable_op_kwargs = [{}]
+        elif isinstance(variable_op_kwargs, dict):
+            variable_op_kwargs = [variable_op_kwargs]
 
+        # construct physical operators for each set of kwargs
         physical_expressions = []
-        for model in physical_op_params["available_models"]:
-            # skip this model if:
-            # 1. this is a pure vision model and we're not doing an image conversion, or
-            # 2. this is a pure text model and we're doing an image conversion, or
-            # 3. this is a vision model hosted by Together (i.e. LLAMA3 vision) and there is more than one image field
-            first_criteria = model in pure_vision_models and not is_image_conversion
-            second_criteria = model in pure_text_models and is_image_conversion
-            third_criteria = model.is_llama_model() and model.is_vision_model() and (num_image_fields > 1 or list_image_field)
-            fourth_criteria = model.is_embedding_model()
-            if first_criteria or second_criteria or third_criteria or fourth_criteria:
-                continue
+        for var_op_kwargs in variable_op_kwargs:
+            # get kwargs for this physical operator instance
+            op_kwargs = {**fixed_op_kwargs, **var_op_kwargs}
 
-            # construct multi-expression
-            op = LLMConvertBonded(
-                model=model,
-                prompt_strategy=PromptStrategy.COT_QA_IMAGE if is_image_conversion else PromptStrategy.COT_QA,
-                **op_kwargs,
-            )
-            expression = PhysicalExpression(
-                operator=op,
-                input_group_ids=logical_expression.input_group_ids,
-                input_fields=logical_expression.input_fields,
-                depends_on_field_names=logical_expression.depends_on_field_names,
-                generated_fields=logical_expression.generated_fields,
-                group_id=logical_expression.group_id,
-            )
-            physical_expressions.append(expression)
+            # construct the physical operator
+            op = physical_op_class(**op_kwargs)
 
-        deduped_physical_expressions = set(physical_expressions)
-        logger.debug(f"Done substituting LLMConvertBondedRule for {logical_expression}")
+            # construct physical expression and add to list of expressions
+            expression = PhysicalExpression.from_op_and_logical_expr(op, logical_expression)
+            physical_expressions.append(expression)
 
-        return deduped_physical_expressions
+        return set(physical_expressions)
 
 
-class CodeSynthesisConvertRule(ImplementationRule):
+class NonLLMConvertRule(ImplementationRule):
     """
-    Base rule for code synthesis convert operators; the physical convert class
-    (CodeSynthesisConvertSingle) is provided by sub-class rules.
-
-    NOTE: we provide the physical convert class(es) in their own sub-classed rules to make
-    it easier to allow/disallow groups of rules at the Optimizer level.
+    Substitute a logical expression for a UDF ConvertScan with a NonLLMConvert physical implementation.
     """
 
-    physical_convert_class = None  # overriden by sub-classes
-
     @classmethod
     def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
-        logical_op = logical_expression.operator
-        is_image_conversion = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        is_match = (
-            isinstance(logical_op, ConvertScan)
-            and not is_image_conversion
-            and logical_op.cardinality != Cardinality.ONE_TO_MANY
-            and logical_op.udf is None
-        )
-        logger.debug(f"CodeSynthesisConvertRule matches_pattern: {is_match} for {logical_expression}")
+        is_match = isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is not None
+        logger.debug(f"NonLLMConvertRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
-        logger.debug(f"Substituting CodeSynthesisConvertRule for {logical_expression}")
-
-        logical_op = logical_expression.operator
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
+        logger.debug(f"Substituting NonLLMConvertRule for {logical_expression}")
+        return cls._perform_substitution(logical_expression, NonLLMConvert, runtime_kwargs)
 
-        # get initial set of parameters for physical op
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
-            {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
-            }
-        )
 
-        # construct multi-expression
-        op = cls.physical_convert_class(
-            exemplar_generation_model=physical_op_params["champion_model"],
-            code_synth_model=physical_op_params["code_champion_model"],
-            fallback_model=physical_op_params["fallback_model"],
-            prompt_strategy=PromptStrategy.COT_QA,
-            **op_kwargs,
-        )
-        expression = PhysicalExpression(
-            operator=op,
-            input_group_ids=logical_expression.input_group_ids,
-            input_fields=logical_expression.input_fields,
-            depends_on_field_names=logical_expression.depends_on_field_names,
-            generated_fields=logical_expression.generated_fields,
-            group_id=logical_expression.group_id,
-        )
-        deduped_physical_expressions = set([expression])
-        logger.debug(f"Done substituting CodeSynthesisConvertRule for {logical_expression}")
+class LLMConvertBondedRule(ImplementationRule):
+    """
+    Substitute a logical expression for a ConvertScan with a bonded convert physical implementation.
+    """
 
-        return deduped_physical_expressions
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
+        is_match = isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is None
+        logger.debug(f"LLMConvertBondedRule matches_pattern: {is_match} for {logical_expression}")
+        return is_match
 
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
+        logger.debug(f"Substituting LLMConvertBondedRule for {logical_expression}")
 
-class CodeSynthesisConvertSingleRule(CodeSynthesisConvertRule):
-    """
-    Substitute a logical expression for a ConvertScan with a (single) code synthesis physical implementation.
-    """
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        # NOTE: right now we exclusively allow image or audio operations, but not both simultaneously
+        prompt_strategy, no_reasoning_prompt_strategy = None, None
+        no_reasoning = runtime_kwargs["reasoning_effort"] in [None, "minimal", "low"]
+        if cls._is_text_only_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_QA
+            no_reasoning_prompt_strategy = PromptStrategy.COT_QA_NO_REASONING
+        elif cls._is_image_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_QA_IMAGE
+            no_reasoning_prompt_strategy = PromptStrategy.COT_QA_IMAGE_NO_REASONING
+        elif cls._is_audio_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_QA_AUDIO
+            no_reasoning_prompt_strategy = PromptStrategy.COT_QA_AUDIO_NO_REASONING
+        variable_op_kwargs = [
+            {
+                "model": model,
+                "prompt_strategy": no_reasoning_prompt_strategy if model.is_reasoning_model() and no_reasoning else prompt_strategy,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
+            }
+            for model in models
+        ]
 
-    physical_convert_class = CodeSynthesisConvertSingle
+        return cls._perform_substitution(logical_expression, LLMConvertBonded, runtime_kwargs, variable_op_kwargs)
 
 
 class RAGConvertRule(ImplementationRule):
@@ -436,68 +531,30 @@ class RAGConvertRule(ImplementationRule):
     @classmethod
     def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         logical_op = logical_expression.operator
-        is_image_conversion = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        is_match = isinstance(logical_op, ConvertScan) and not is_image_conversion and logical_op.udf is None
+        is_match = isinstance(logical_op, ConvertScan) and cls._is_text_only_operation(logical_expression) and logical_op.udf is None
         logger.debug(f"RAGConvertRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting RAGConvertRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-
-        # get initial set of parameters for physical op
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        variable_op_kwargs = [
             {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
+                "model": model,
+                "prompt_strategy": PromptStrategy.COT_QA,
+                "num_chunks_per_field": num_chunks_per_field,
+                "chunk_size": chunk_size,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
             }
-        )
+            for model in models
+            for num_chunks_per_field in cls.num_chunks_per_fields
+            for chunk_size in cls.chunk_sizes
+        ]
 
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-        pure_vision_models = {model for model in vision_models if model not in text_models}
-
-        physical_expressions = []
-        for model in physical_op_params["available_models"]:
-            # skip this model if this is a pure image model
-            if model in pure_vision_models or model.is_embedding_model():
-                continue
-
-            for num_chunks_per_field in cls.num_chunks_per_fields:
-                for chunk_size in cls.chunk_sizes:
-                    # construct multi-expression
-                    op = RAGConvert(
-                        model=model,
-                        prompt_strategy=PromptStrategy.COT_QA,
-                        num_chunks_per_field=num_chunks_per_field,
-                        chunk_size=chunk_size,
-                        **op_kwargs,
-                    )
-                    expression = PhysicalExpression(
-                        operator=op,
-                        input_group_ids=logical_expression.input_group_ids,
-                        input_fields=logical_expression.input_fields,
-                        depends_on_field_names=logical_expression.depends_on_field_names,
-                        generated_fields=logical_expression.generated_fields,
-                        group_id=logical_expression.group_id,
-                    )
-                    physical_expressions.append(expression)
-
-        logger.debug(f"Done substituting RAGConvertRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
-
-        return deduped_physical_expressions
+        return cls._perform_substitution(logical_expression, RAGConvert, runtime_kwargs, variable_op_kwargs)
 
 
 class MixtureOfAgentsConvertRule(ImplementationRule):
@@ -511,93 +568,35 @@ class MixtureOfAgentsConvertRule(ImplementationRule):
     @classmethod
     def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         logical_op = logical_expression.operator
-        is_match = isinstance(logical_op, ConvertScan) and logical_op.udf is None
+        # TODO: remove audio limitation once I add prompts
+        is_match = isinstance(logical_op, ConvertScan) and logical_op.udf is None and not cls._is_audio_operation(logical_expression)
         logger.debug(f"MixtureOfAgentsConvertRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting MixtureOfAgentsConvertRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-
-        # get initial set of parameters for physical op
-        op_kwargs: dict = logical_op.get_logical_op_params()
-        op_kwargs.update(
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        proposer_model_set = {model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)}
+        aggregator_model_set = {model for model in runtime_kwargs["available_models"] if model.is_text_model()}
+        proposer_prompt_strategy = PromptStrategy.COT_MOA_PROPOSER_IMAGE if cls._is_image_operation(logical_expression) else PromptStrategy.COT_MOA_PROPOSER
+        variable_op_kwargs = [
             {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
+                "proposer_models": list(proposer_models),
+                "temperatures": [temp] * len(proposer_models),
+                "aggregator_model": aggregator_model,
+                "proposer_prompt_strategy": proposer_prompt_strategy,
+                "aggregator_prompt_strategy": PromptStrategy.COT_MOA_AGG,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
             }
-        )
+            for k in cls.num_proposer_models
+            for temp in cls.temperatures
+            for proposer_models in combinations(proposer_model_set, k)
+            for aggregator_model in aggregator_model_set
+        ]
 
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-
-        # construct set of proposer models and set of aggregator models
-        num_image_fields = sum(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        list_image_field = any(
-            [
-                field.is_image_field and hasattr(field, "element_type")
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        proposer_model_set, is_image_conversion = text_models, False
-        if num_image_fields > 1 or list_image_field:
-            proposer_model_set = [model for model in vision_models if not model.is_llama_model()]
-            is_image_conversion = True
-        elif num_image_fields == 1:
-            proposer_model_set = vision_models
-            is_image_conversion = True
-        aggregator_model_set = text_models
-
-        # filter un-available models out of sets
-        proposer_model_set = {model for model in proposer_model_set if model in physical_op_params["available_models"]}
-        aggregator_model_set = {
-            model for model in aggregator_model_set if model in physical_op_params["available_models"]
-        }
-
-        # construct MixtureOfAgentsConvert operations for various numbers of proposer models
-        # and for every combination of proposer models and aggregator model
-        physical_expressions = []
-        for k in cls.num_proposer_models:
-            for temp in cls.temperatures:
-                for proposer_models in combinations(proposer_model_set, k):
-                    for aggregator_model in aggregator_model_set:
-                        # construct multi-expression
-                        op = MixtureOfAgentsConvert(
-                            proposer_models=list(proposer_models),
-                            temperatures=[temp] * len(proposer_models),
-                            aggregator_model=aggregator_model,
-                            proposer_prompt=op_kwargs.get("prompt"),
-                            proposer_prompt_strategy=PromptStrategy.COT_MOA_PROPOSER_IMAGE
-                            if is_image_conversion
-                            else PromptStrategy.COT_MOA_PROPOSER,
-                            aggregator_prompt_strategy=PromptStrategy.COT_MOA_AGG,
-                            **op_kwargs,
-                        )
-                        expression = PhysicalExpression(
-                            operator=op,
-                            input_group_ids=logical_expression.input_group_ids,
-                            input_fields=logical_expression.input_fields,
-                            depends_on_field_names=logical_expression.depends_on_field_names,
-                            generated_fields=logical_expression.generated_fields,
-                            group_id=logical_expression.group_id,
-                        )
-                        physical_expressions.append(expression)
-
-        logger.debug(f"Done substituting MixtureOfAgentsConvertRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
-
-        return deduped_physical_expressions
+        return cls._perform_substitution(logical_expression, MixtureOfAgentsConvert, runtime_kwargs, variable_op_kwargs)
 
 
 class CriticAndRefineConvertRule(ImplementationRule):
@@ -608,99 +607,32 @@ class CriticAndRefineConvertRule(ImplementationRule):
     @classmethod
     def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         logical_op = logical_expression.operator
-        is_match = isinstance(logical_op, ConvertScan) and logical_op.udf is None
+        # TODO: remove audio limitation once I add prompts
+        is_match = isinstance(logical_op, ConvertScan) and logical_op.udf is None and not cls._is_audio_operation(logical_expression)
         logger.debug(f"CriticAndRefineConvertRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting CriticAndRefineConvertRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-
-        # Get initial parameters for physical operator
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        prompt_strategy = PromptStrategy.COT_QA_IMAGE if cls._is_image_operation(logical_expression) else PromptStrategy.COT_QA
+        variable_op_kwargs = [
             {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
+                "model": model,
+                "critic_model": critic_model,
+                "refine_model": refine_model,
+                "prompt_strategy": prompt_strategy,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
             }
-        )
+            for model in models
+            for critic_model in models
+            for refine_model in models
+        ]
 
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-        pure_text_models = {model for model in text_models if model not in vision_models}
-        pure_vision_models = {model for model in vision_models if model not in text_models}
-
-        # compute attributes about this convert operation
-        is_image_conversion = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        num_image_fields = sum(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        list_image_field = any(
-            [
-                field.is_image_field and hasattr(field, "element_type")
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-
-        # identify models which can be used for this convert operation
-        models = []
-        for model in physical_op_params["available_models"]:
-            # skip this model if:
-            # 1. this is a pure vision model and we're not doing an image conversion, or
-            # 2. this is a pure text model and we're doing an image conversion, or
-            # 3. this is a vision model hosted by Together (i.e. LLAMA3 vision) and there is more than one image field
-            first_criteria = model in pure_vision_models and not is_image_conversion
-            second_criteria = model in pure_text_models and is_image_conversion
-            third_criteria = model.is_llama_model() and model.is_vision_model() and (num_image_fields > 1 or list_image_field)
-            fourth_criteria = model.is_embedding_model()
-            if first_criteria or second_criteria or third_criteria or fourth_criteria:
-                continue
-
-            models.append(model)
-
-        # TODO: heuristic(s) to narrow the space of critic and refine models we consider using class attributes
-        # construct CriticAndRefineConvert operations for every combination of model, critic model, and refinement model
-        physical_expressions = []
-        for model in models:
-            for critic_model in models:
-                for refine_model in models:
-                    # construct multi-expression
-                    op = CriticAndRefineConvert(
-                        model=model,
-                        prompt_strategy=PromptStrategy.COT_QA_IMAGE if is_image_conversion else PromptStrategy.COT_QA,
-                        critic_model=critic_model,
-                        refine_model=refine_model,
-                        **op_kwargs,
-                    )
-                    expression = PhysicalExpression(
-                        operator=op,
-                        input_group_ids=logical_expression.input_group_ids,
-                        input_fields=logical_expression.input_fields,
-                        depends_on_field_names=logical_expression.depends_on_field_names,
-                        generated_fields=logical_expression.generated_fields,
-                        group_id=logical_expression.group_id,
-                    )
-                    physical_expressions.append(expression)
-
-        logger.debug(f"Done substituting CriticAndRefineConvertRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
-
-        return deduped_physical_expressions
+        return cls._perform_substitution(logical_expression, CriticAndRefineConvert, runtime_kwargs, variable_op_kwargs)
 
 
 class SplitConvertRule(ImplementationRule):
@@ -713,67 +645,29 @@ class SplitConvertRule(ImplementationRule):
     @classmethod
     def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         logical_op = logical_expression.operator
-        is_image_conversion = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        is_match = isinstance(logical_op, ConvertScan) and not is_image_conversion and logical_op.udf is None
+        is_match = isinstance(logical_op, ConvertScan) and cls._is_text_only_operation() and logical_op.udf is None
         logger.debug(f"SplitConvertRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting SplitConvertRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-
-        # get initial set of parameters for physical op
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        variable_op_kwargs = [
             {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
+                "model": model,
+                "min_size_to_chunk": min_size_to_chunk,
+                "num_chunks": num_chunks,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
             }
-        )
-
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-        pure_vision_models = {model for model in vision_models if model not in text_models}
-
-        physical_expressions = []
-        for model in physical_op_params["available_models"]:
-            # skip this model if this is a pure image model
-            if model in pure_vision_models or model.is_embedding_model():
-                continue
+            for model in models
+            for min_size_to_chunk in cls.min_size_to_chunk
+            for num_chunks in cls.num_chunks
+        ]
 
-            for min_size_to_chunk in cls.min_size_to_chunk:
-                for num_chunks in cls.num_chunks:
-                    # construct multi-expression
-                    op = SplitConvert(
-                        model=model,
-                        num_chunks=num_chunks,
-                        min_size_to_chunk=min_size_to_chunk,
-                        **op_kwargs,
-                    )
-                    expression = PhysicalExpression(
-                        operator=op,
-                        input_group_ids=logical_expression.input_group_ids,
-                        input_fields=logical_expression.input_fields,
-                        depends_on_field_names=logical_expression.depends_on_field_names,
-                        generated_fields=logical_expression.generated_fields,
-                        group_id=logical_expression.group_id,
-                    )
-                    physical_expressions.append(expression)
-
-        logger.debug(f"Done substituting SplitConvertRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
-
-        return deduped_physical_expressions
+        return cls._perform_substitution(logical_expression, SplitConvert, runtime_kwargs, variable_op_kwargs)
 
 
 class RetrieveRule(ImplementationRule):
@@ -789,42 +683,13 @@ class RetrieveRule(ImplementationRule):
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting RetrieveRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-
-        physical_expressions = []
-        ks = cls.k_budgets if logical_op.k == -1 else [logical_op.k]
-        for k in ks:
-            # get initial set of parameters for physical op
-            op_kwargs = logical_op.get_logical_op_params()
-            op_kwargs.update(
-                {
-                    "verbose": physical_op_params["verbose"],
-                    "logical_op_id": logical_op.get_logical_op_id(),
-                    "logical_op_name": logical_op.logical_op_name(),
-                    "k": k,
-                }
-            )
-
-            # construct multi-expression
-            op = RetrieveOp(**op_kwargs)
-            expression = PhysicalExpression(
-                operator=op,
-                input_group_ids=logical_expression.input_group_ids,
-                input_fields=logical_expression.input_fields,
-                depends_on_field_names=logical_expression.depends_on_field_names,
-                generated_fields=logical_expression.generated_fields,
-                group_id=logical_expression.group_id,
-            )
-
-            physical_expressions.append(expression)
-
-        logger.debug(f"Done substituting RetrieveRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
-
-        return deduped_physical_expressions
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        ks = cls.k_budgets if logical_expression.operator.k == -1 else [logical_expression.operator.k]
+        variable_op_kwargs = [{"k": k} for k in ks]
+        return cls._perform_substitution(logical_expression, RetrieveOp, runtime_kwargs, variable_op_kwargs)
 
 
 class NonLLMFilterRule(ImplementationRule):
@@ -832,8 +697,8 @@ class NonLLMFilterRule(ImplementationRule):
     Substitute a logical expression for a FilteredScan with a non-llm filter physical implementation.
     """
 
-    @staticmethod
-    def matches_pattern(logical_expression: LogicalExpression) -> bool:
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         is_match = (
             isinstance(logical_expression.operator, FilteredScan)
             and logical_expression.operator.filter.filter_fn is not None
@@ -841,33 +706,10 @@ class NonLLMFilterRule(ImplementationRule):
         logger.debug(f"NonLLMFilterRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
-    @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting NonLLMFilterRule for {logical_expression}")
-
-        logical_op = logical_expression.operator
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
-            {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
-            }
-        )
-        op = NonLLMFilter(**op_kwargs)
-
-        expression = PhysicalExpression(
-            operator=op,
-            input_group_ids=logical_expression.input_group_ids,
-            input_fields=logical_expression.input_fields,
-            depends_on_field_names=logical_expression.depends_on_field_names,
-            generated_fields=logical_expression.generated_fields,
-            group_id=logical_expression.group_id,
-        )
-        logger.debug(f"Done substituting NonLLMFilterRule for {logical_expression}")
-        deduped_physical_expressions = set([expression])
-
-        return deduped_physical_expressions
+        return cls._perform_substitution(logical_expression, NonLLMFilter, runtime_kwargs)
 
 
 class LLMFilterRule(ImplementationRule):
@@ -875,8 +717,8 @@ class LLMFilterRule(ImplementationRule):
     Substitute a logical expression for a FilteredScan with an llm filter physical implementation.
     """
 
-    @staticmethod
-    def matches_pattern(logical_expression: LogicalExpression) -> bool:
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         is_match = (
             isinstance(logical_expression.operator, FilteredScan)
             and logical_expression.operator.filter.filter_condition is not None
@@ -884,82 +726,76 @@ class LLMFilterRule(ImplementationRule):
         logger.debug(f"LLMFilterRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
-    @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting LLMFilterRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        # NOTE: right now we exclusively allow image or audio operations, but not both simultaneously
+        prompt_strategy, no_reasoning_prompt_strategy = None, None
+        no_reasoning = runtime_kwargs["reasoning_effort"] in [None, "minimal", "low"]
+        if cls._is_text_only_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_BOOL
+            no_reasoning_prompt_strategy = PromptStrategy.COT_BOOL_NO_REASONING
+        elif cls._is_image_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_BOOL_IMAGE
+            no_reasoning_prompt_strategy = PromptStrategy.COT_BOOL_IMAGE_NO_REASONING
+        elif cls._is_audio_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_BOOL_AUDIO
+            no_reasoning_prompt_strategy = PromptStrategy.COT_BOOL_AUDIO_NO_REASONING
+        variable_op_kwargs = [
             {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
+                "model": model,
+                "prompt_strategy": no_reasoning_prompt_strategy if model.is_reasoning_model() and no_reasoning else prompt_strategy,
+                "reasoning_effort": runtime_kwargs["reasoning_effort"]
             }
-        )
+            for model in models
+        ]
 
-        # identify models which can be used strictly for text or strictly for images
-        vision_models = set(get_vision_models())
-        text_models = set(get_models())
-        pure_text_models = {model for model in text_models if model not in vision_models}
-        pure_vision_models = {model for model in vision_models if model not in text_models}
-
-        # compute attributes about this filter operation
-        is_image_filter = any(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        num_image_fields = sum(
-            [
-                field.is_image_field
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
-        list_image_field = any(
-            [
-                field.is_image_field and hasattr(field, "element_type")
-                for field_name, field in logical_expression.input_fields.items()
-                if field_name.split(".")[-1] in logical_expression.depends_on_field_names
-            ]
-        )
+        return cls._perform_substitution(logical_expression, LLMFilter, runtime_kwargs, variable_op_kwargs)
 
-        physical_expressions = []
-        for model in physical_op_params["available_models"]:
-            # skip this model if:
-            # 1. this is a pure vision model and we're not doing an image filter, or
-            # 2. this is a pure text model and we're doing an image filter, or
-            # 3. this is a vision model hosted by Together (i.e. LLAMA3 vision) and there is more than one image field
-            first_criteria = model in pure_vision_models and not is_image_filter
-            second_criteria = model in pure_text_models and is_image_filter
-            third_criteria = model.is_llama_model() and model.is_vision_model() and (num_image_fields > 1 or list_image_field)
-            fourth_criteria = model.is_embedding_model()
-            if first_criteria or second_criteria or third_criteria or fourth_criteria:
-                continue
 
-            # construct multi-expression
-            op = LLMFilter(
-                model=model,
-                prompt_strategy=PromptStrategy.COT_BOOL_IMAGE if is_image_filter else PromptStrategy.COT_BOOL,
-                **op_kwargs,
-            )
-            expression = PhysicalExpression(
-                operator=op,
-                input_group_ids=logical_expression.input_group_ids,
-                input_fields=logical_expression.input_fields,
-                depends_on_field_names=logical_expression.depends_on_field_names,
-                generated_fields=logical_expression.generated_fields,
-                group_id=logical_expression.group_id,
-            )
-            physical_expressions.append(expression)
+class LLMJoinRule(ImplementationRule):
+    """
+    Substitute a logical expression for a JoinOp with an (LLM) NestedLoopsJoin physical implementation.
+    """
 
-        logger.debug(f"Done substituting LLMFilterRule for {logical_expression}")
-        deduped_physical_expressions = set(physical_expressions)
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
+        is_match = isinstance(logical_expression.operator, JoinOp)
+        logger.debug(f"LLMJoinRule matches_pattern: {is_match} for {logical_expression}")
+        return is_match
 
-        return deduped_physical_expressions
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
+        logger.debug(f"Substituting LLMJoinRule for {logical_expression}")
+
+        # create variable physical operator kwargs for each model which can implement this logical_expression
+        models = [model for model in runtime_kwargs["available_models"] if cls._model_matches_input(model, logical_expression)]
+        # NOTE: right now we exclusively allow image or audio operations, but not both simultaneously
+        prompt_strategy, no_reasoning_prompt_strategy = None, None
+        no_reasoning = runtime_kwargs["reasoning_effort"] in [None, "minimal", "low"]
+        if cls._is_text_only_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_JOIN
+            no_reasoning_prompt_strategy = PromptStrategy.COT_JOIN_NO_REASONING
+        elif cls._is_image_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_JOIN_IMAGE
+            no_reasoning_prompt_strategy = PromptStrategy.COT_JOIN_IMAGE_NO_REASONING
+        elif cls._is_audio_operation(logical_expression):
+            prompt_strategy = PromptStrategy.COT_JOIN_AUDIO
+            no_reasoning_prompt_strategy = PromptStrategy.COT_JOIN_AUDIO_NO_REASONING
+        variable_op_kwargs = [
+            {
+                "model": model,
+                "prompt_strategy": no_reasoning_prompt_strategy if model.is_reasoning_model() and no_reasoning else prompt_strategy,
+                "join_parallelism": runtime_kwargs["join_parallelism"],
+                "reasoning_effort": runtime_kwargs["reasoning_effort"],
+            }
+            for model in models
+        ]
+
+        return cls._perform_substitution(logical_expression, NestedLoopsJoin, runtime_kwargs, variable_op_kwargs)
 
 
 class AggregateRule(ImplementationRule):
@@ -967,47 +803,71 @@ class AggregateRule(ImplementationRule):
     Substitute the logical expression for an aggregate with its physical counterpart.
     """
 
-    @staticmethod
-    def matches_pattern(logical_expression: LogicalExpression) -> bool:
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
         is_match = isinstance(logical_expression.operator, Aggregate)
         logger.debug(f"AggregateRule matches_pattern: {is_match} for {logical_expression}")
         return is_match
 
-    @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting AggregateRule for {logical_expression}")
 
-        logical_op = logical_expression.operator
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
-            {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
-            }
-        )
-
-        op = None
-        if logical_op.agg_func == AggFunc.COUNT:
-            op = CountAggregateOp(**op_kwargs)
-        elif logical_op.agg_func == AggFunc.AVERAGE:
-            op = AverageAggregateOp(**op_kwargs)
+        # get the physical op class based on the aggregation function
+        physical_op_class = None
+        if logical_expression.operator.agg_func == AggFunc.COUNT:
+            physical_op_class = CountAggregateOp
+        elif logical_expression.operator.agg_func == AggFunc.AVERAGE:
+            physical_op_class = AverageAggregateOp
         else:
-            raise Exception(f"Cannot support aggregate function: {logical_op.agg_func}")
-
-        expression = PhysicalExpression(
-            operator=op,
-            input_group_ids=logical_expression.input_group_ids,
-            input_fields=logical_expression.input_fields,
-            depends_on_field_names=logical_expression.depends_on_field_names,
-            generated_fields=logical_expression.generated_fields,
-            group_id=logical_expression.group_id,
-        )
+            raise Exception(f"Cannot support aggregate function: {logical_expression.operator.agg_func}")
 
-        logger.debug(f"Done substituting AggregateRule for {logical_expression}")
-        deduped_physical_expressions = set([expression])
+        # perform the substitution
+        return cls._perform_substitution(logical_expression, physical_op_class, runtime_kwargs)
+
+
+class AddContextsBeforeComputeRule(ImplementationRule):
+    """
+    Searches the ContextManager for additional contexts which may be useful for the given computation.
+
+    TODO: track cost of generating search query
+    """
+    k = 1
+    SEARCH_GENERATOR_PROMPT = CONTEXT_SEARCH_PROMPT
+
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
+        is_match = isinstance(logical_expression.operator, ComputeOperator)
+        logger.debug(f"AddContextsBeforeComputeRule matches_pattern: {is_match} for {logical_expression}")
+        return is_match
 
-        return deduped_physical_expressions
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
+        logger.debug(f"Substituting AddContextsBeforeComputeRule for {logical_expression}")
+
+        # load an LLM to generate a short search query
+        model = None
+        if os.getenv("OPENAI_API_KEY"):
+            model = "openai/gpt-4o-mini"
+        elif os.getenv("ANTHROPIC_API_KEY"):
+            model = "anthropic/claude-3-5-sonnet-20241022"
+        elif os.getenv("GEMINI_API_KEY"):
+            model = "vertex_ai/gemini-2.0-flash"
+        elif os.getenv("TOGETHER_API_KEY"):
+            model = "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"
+
+        # importing litellm here because importing above causes deprecation warning
+        import litellm
+
+        # retrieve any additional context which may be useful
+        cm = ContextManager()
+        response = litellm.completion(
+            model=model,
+            messages=[{"role": "user", "content": cls.SEARCH_GENERATOR_PROMPT.format(instruction=logical_expression.operator.instruction)}]
+        )
+        query = response.choices[0].message.content
+        variable_op_kwargs = {"additional_contexts": cm.search_context(query, k=cls.k, where={"materialized": True})}
+        return cls._perform_substitution(logical_expression, SmolAgentsCompute, runtime_kwargs, variable_op_kwargs)
 
 
 class BasicSubstitutionRule(ImplementationRule):
@@ -1018,11 +878,13 @@ class BasicSubstitutionRule(ImplementationRule):
 
     LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP = {
         BaseScan: MarshalAndScanDataOp,
-        CacheScan: CacheScanDataOp,
+        # ComputeOperator: SmolAgentsCompute,
+        SearchOperator: SmolAgentsSearch, # SmolAgentsManagedSearch, # SmolAgentsCustomManagedSearch
+        ContextScan: ContextScanOp,
+        Distinct: DistinctOp,
         LimitScan: LimitScanOp,
         Project: ProjectOp,
         GroupByAggregate: ApplyGroupByOp,
-        MapScan: MapOp,
     }
 
     @classmethod
@@ -1033,31 +895,7 @@ class BasicSubstitutionRule(ImplementationRule):
         return is_match
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **runtime_kwargs) -> set[PhysicalExpression]:
         logger.debug(f"Substituting BasicSubstitutionRule for {logical_expression}")
-
-        logical_op = logical_expression.operator
-        op_kwargs = logical_op.get_logical_op_params()
-        op_kwargs.update(
-            {
-                "verbose": physical_op_params["verbose"],
-                "logical_op_id": logical_op.get_logical_op_id(),
-                "logical_op_name": logical_op.logical_op_name(),
-            }
-        )
-        physical_op_class = cls.LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP[logical_op.__class__]
-        op = physical_op_class(**op_kwargs)
-
-        expression = PhysicalExpression(
-            operator=op,
-            input_group_ids=logical_expression.input_group_ids,
-            input_fields=logical_expression.input_fields,
-            depends_on_field_names=logical_expression.depends_on_field_names,
-            generated_fields=logical_expression.generated_fields,
-            group_id=logical_expression.group_id,
-        )
-
-        logger.debug(f"Done substituting BasicSubstitutionRule for {logical_expression}")
-        deduped_physical_expressions = set([expression])
-
-        return deduped_physical_expressions
+        physical_op_class = cls.LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP[logical_expression.operator.__class__]
+        return cls._perform_substitution(logical_expression, physical_op_class, runtime_kwargs)