palimpzest 0.5.3__py3-none-any.whl → 0.6.0__py3-none-any.whl

palimpzest/query/optimizer/rules.py

@@ -1,13 +1,11 @@
 from copy import deepcopy
 from itertools import combinations
-from typing import Dict, Set, Tuple
 
 from palimpzest.constants import AggFunc, Cardinality, Model, PromptStrategy
-from palimpzest.core.lib.fields import ListField
 from palimpzest.query.operators.aggregate import ApplyGroupByOp, AverageAggregateOp, CountAggregateOp
 from palimpzest.query.operators.code_synthesis_convert import CodeSynthesisConvertSingle
 from palimpzest.query.operators.convert import LLMConvertBonded, LLMConvertConventional, NonLLMConvert
-from palimpzest.query.operators.datasource import CacheScanDataOp, MarshalAndScanDataOp
+from palimpzest.query.operators.critique_and_refine_convert import CriticAndRefineConvert
 from palimpzest.query.operators.filter import LLMFilter, NonLLMFilter
 from palimpzest.query.operators.limit import LimitScanOp
 from palimpzest.query.operators.logical import (
@@ -25,6 +23,7 @@ from palimpzest.query.operators.mixture_of_agents_convert import MixtureOfAgents
 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.token_reduction_convert import (
     TokenReducedConvertBonded,
     TokenReducedConvertConventional,
@@ -47,7 +46,7 @@ class Rule:
         raise NotImplementedError("Calling this method from an abstract base class.")
 
     @staticmethod
-    def substitute(logical_expression: LogicalExpression, **kwargs) -> Set[Expression]:
+    def substitute(logical_expression: LogicalExpression, **kwargs) -> set[Expression]:
         raise NotImplementedError("Calling this method from an abstract base class.")
 
 
@@ -60,8 +59,8 @@ class TransformationRule(Rule):
 
     @staticmethod
     def substitute(
-        logical_expression: LogicalExpression, groups: Dict[int, Group], expressions: Dict[int, Expression], **kwargs
-    ) -> Tuple[Set[LogicalExpression], Set[Group]]:
+        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
         potentially creates new intermediate expressions and groups.
@@ -86,8 +85,8 @@ class PushDownFilter(TransformationRule):
 
     @staticmethod
     def substitute(
-        logical_expression: LogicalExpression, groups: Dict[int, Group], expressions: Dict[int, Expression], **kwargs
-    ) -> Tuple[Set[LogicalExpression], Set[Group]]:
+        logical_expression: LogicalExpression, groups: dict[int, Group], expressions: dict[int, Expression], **kwargs
+    ) -> tuple[set[LogicalExpression], set[Group]]:
         # initialize the sets of new logical expressions and groups to be returned
         new_logical_expressions, new_groups = set(), set()
 
@@ -103,7 +102,7 @@ class PushDownFilter(TransformationRule):
                 continue
 
             # iterate over logical expressions
-            logical_exprs = input_group.logical_expressions.copy()
+            logical_exprs = deepcopy(input_group.logical_expressions)
             for expr in logical_exprs:
                 # if the expression operator is not a convert or a filter, we cannot swap
                 if not (isinstance(expr.operator, (ConvertScan, FilteredScan))):
@@ -114,10 +113,10 @@ class PushDownFilter(TransformationRule):
                     continue
 
                 # create new logical expression with filter pushed down to the input group's logical expression
-                new_input_group_ids = expr.input_group_ids.copy()
-                new_input_fields = expr.input_fields.copy()
-                new_depends_on_field_names = logical_expression.depends_on_field_names.copy()
-                new_generated_fields = logical_expression.generated_fields.copy()
+                new_input_group_ids = deepcopy(expr.input_group_ids)
+                new_input_fields = deepcopy(expr.input_fields)
+                new_depends_on_field_names = deepcopy(logical_expression.depends_on_field_names)
+                new_generated_fields = deepcopy(logical_expression.generated_fields)
                 new_filter_expr = LogicalExpression(
                     filter_operator,
                     input_group_ids=new_input_group_ids,
@@ -215,7 +214,7 @@ class NonLLMConvertRule(ImplementationRule):
         return isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is not None
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -259,7 +258,7 @@ class LLMConvertRule(ImplementationRule):
         return isinstance(logical_expression.operator, ConvertScan) and logical_expression.operator.udf is None
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -293,7 +292,7 @@ class LLMConvertRule(ImplementationRule):
             if field_name.split(".")[-1] in logical_expression.depends_on_field_names
         ])
         list_image_field = any([
-            field.is_image_field and isinstance(field, ListField)
+            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
         ])
@@ -368,7 +367,7 @@ class TokenReducedConvertRule(ImplementationRule):
         return isinstance(logical_op, ConvertScan) and not is_image_conversion and logical_op.udf is None
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -459,7 +458,7 @@ class CodeSynthesisConvertRule(ImplementationRule):
         )
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -518,7 +517,7 @@ class RAGConvertRule(ImplementationRule):
         return isinstance(logical_op, ConvertScan) and not is_image_conversion and logical_op.udf is None
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -580,7 +579,7 @@ class MixtureOfAgentsConvertRule(ImplementationRule):
         return isinstance(logical_op, ConvertScan) and logical_op.udf is None
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         # get initial set of parameters for physical op
@@ -605,7 +604,7 @@ class MixtureOfAgentsConvertRule(ImplementationRule):
             if field_name.split(".")[-1] in logical_expression.depends_on_field_names
         ])
         list_image_field = any([
-            field.is_image_field and isinstance(field, ListField)
+            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
         ])
@@ -651,6 +650,98 @@ class MixtureOfAgentsConvertRule(ImplementationRule):
 
         return set(physical_expressions)
 
+class CriticAndRefineConvertRule(ImplementationRule):
+    """
+    Implementation rule for the CriticAndRefineConvert operator.
+    """
+
+    @classmethod
+    def matches_pattern(cls, logical_expression: LogicalExpression) -> bool:
+        logical_op = logical_expression.operator
+        return isinstance(logical_op, ConvertScan) and logical_op.udf is None
+
+    @classmethod
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
+        logical_op = logical_expression.operator
+
+        # Get initial parameters for physical 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(),
+            }
+        )
+
+        # NOTE: when comparing pz.Model(s), equality is determined by the string (i.e. pz.Model.value)
+        #       thus, Model.GPT_4o and Model.GPT_4o_V map to the same value; this allows us to use set logic
+        #
+        # 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_V) 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 == Model.LLAMA3_V and (num_image_fields > 1 or list_image_field)
+            if first_criteria or second_criteria or third_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)
+
+        # Return the set containing the new physical expression
+        return set(physical_expressions)
+
 
 class RetrieveRule(ImplementationRule):
     """
@@ -667,7 +758,7 @@ class RetrieveRule(ImplementationRule):
     @classmethod
     def substitute(
         cls, logical_expression: LogicalExpression, **physical_op_params
-    ) -> Set[PhysicalExpression]:
+    ) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
 
         physical_expressions = []
@@ -713,7 +804,7 @@ class NonLLMFilterRule(ImplementationRule):
         )
 
     @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
         op_kwargs = logical_op.get_logical_op_params()
         op_kwargs.update(
@@ -749,7 +840,7 @@ class LLMFilterRule(ImplementationRule):
         )
 
     @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
         op_kwargs = logical_op.get_logical_op_params()
         op_kwargs.update({
@@ -779,7 +870,7 @@ class LLMFilterRule(ImplementationRule):
             if field_name.split(".")[-1] in logical_expression.depends_on_field_names
         ])
         list_image_field = any([
-            field.is_image_field and isinstance(field, ListField)
+            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
         ])
@@ -825,7 +916,7 @@ class AggregateRule(ImplementationRule):
         return isinstance(logical_expression.operator, Aggregate)
 
     @staticmethod
-    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
         op_kwargs = logical_op.get_logical_op_params()
         op_kwargs.update(
@@ -875,7 +966,7 @@ class BasicSubstitutionRule(ImplementationRule):
         return logical_op_class in cls.LOGICAL_OP_CLASS_TO_PHYSICAL_OP_CLASS_MAP
 
     @classmethod
-    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> Set[PhysicalExpression]:
+    def substitute(cls, logical_expression: LogicalExpression, **physical_op_params) -> set[PhysicalExpression]:
         logical_op = logical_expression.operator
         op_kwargs = logical_op.get_logical_op_params()
         op_kwargs.update(

palimpzest/query/processor/config.py

@@ -2,6 +2,7 @@ import json
 from dataclasses import dataclass, field
 
 from palimpzest.constants import Model
+from palimpzest.core.data.datareaders import DataReader
 from palimpzest.policy import MaxQuality, Policy
 
 
@@ -14,6 +15,8 @@ class QueryProcessorConfig:
     execution_strategy: str = field(default="sequential")
     optimizer_strategy: str = field(default="pareto")
 
+    val_datasource: DataReader | None = field(default=None)
+
     policy: Policy = field(default_factory=MaxQuality)
     scan_start_idx: int = field(default=0)
     num_samples: int = field(default=float("inf"))
@@ -31,8 +34,9 @@ class QueryProcessorConfig:
     allow_model_selection: bool = field(default=True)
     allow_code_synth: bool = field(default=False)
     allow_token_reduction: bool = field(default=False)
-    allow_rag_reduction: bool = field(default=True)
+    allow_rag_reduction: bool = field(default=False)
     allow_mixtures: bool = field(default=True)
+    allow_critic: bool = field(default=False)
     use_final_op_quality: bool = field(default=False)
 
     def to_json_str(self):
@@ -40,6 +44,7 @@ class QueryProcessorConfig:
             "processing_strategy": self.processing_strategy,
             "execution_strategy": self.execution_strategy,
             "optimizer_strategy": self.optimizer_strategy,
+            "val_datasource": None if self.val_datasource is None else self.val_datasource.serialize(),
             "policy": self.policy.to_json_str(),
             "scan_start_idx": self.scan_start_idx,
             "num_samples": self.num_samples,
@@ -57,5 +62,11 @@ class QueryProcessorConfig:
             "allow_token_reduction": self.allow_token_reduction,
             "allow_rag_reduction": self.allow_rag_reduction,
             "allow_mixtures": self.allow_mixtures,
+            "allow_critic": self.allow_critic,
             "use_final_op_quality": self.use_final_op_quality,
         }, indent=2)
+
+    def update(self, **kwargs) -> None:
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)