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

palimpzest/query/generators/generators.py

@@ -101,7 +101,7 @@ def get_json_from_answer(answer: str, model: Model, cardinality: Cardinality) ->
 # TODO: make sure answer parsing works with custom prompts / parsers (can defer this)
 class Generator(Generic[ContextType, InputType]):
     """
-    Abstract base class for Generators.
+    Class for generating new fields for a record using an LLM.
     """
 
     def __init__(
@@ -181,11 +181,11 @@ class Generator(Generic[ContextType, InputType]):
 
         return None
 
-    def _check_bool_answer_text(self, answer_text: str) -> dict | None:
+    def _check_bool_answer_text(self, answer_text: str, throw_exception: bool=False) -> dict | None:
         """
         Return {"passed_operator": True} if and only if "true" is in the answer text.
         Return {"passed_operator": False} if and only if "false" is in the answer text.
-        Otherwise, return None.
+        Otherwise, raise an exception.
         """
         # NOTE: we may be able to eliminate this condition by specifying this JSON output in the prompt;
         # however, that would also need to coincide with a change to allow the parse_answer_fn to set "passed_operator"
@@ -194,6 +194,9 @@ class Generator(Generic[ContextType, InputType]):
         elif "false" in answer_text.lower():
             return {"passed_operator": False}
 
+        if throw_exception:
+            raise Exception(f"Could not parse answer from completion text: {answer_text}")
+
         return None
 
     def _parse_convert_answer(self, completion_text: str, fields: dict[str, FieldInfo], json_output: bool) -> dict[str, list]:
@@ -235,7 +238,7 @@ class Generator(Generic[ContextType, InputType]):
 
         return self._check_convert_answer_text(completion_text, fields, throw_exception=True)
 
-    def _parse_bool_answer(self, completion_text: str) -> dict[str, list]:
+    def _parse_bool_answer(self, completion_text: str, json_output: bool) -> dict[str, list]:
         """Extract the answer from the completion object for filter and join operations."""
         # if the model followed the default instructions, the completion text will place
         # its answer between "ANSWER:" and "---"
@@ -243,6 +246,12 @@ class Generator(Generic[ContextType, InputType]):
         matches = regex.findall(completion_text)
         if len(matches) > 0:
             answer_text = matches[0].strip()
+
+            # if we don't expect a JSON output, return the answer text as is
+            if not json_output:
+                return answer_text
+
+            # otherwise, try to parse the answer text into a JSON object
             field_answers = self._check_bool_answer_text(answer_text)
             if field_answers is not None:
                 return field_answers
@@ -252,16 +261,21 @@ class Generator(Generic[ContextType, InputType]):
         matches = regex.findall(completion_text)
         if len(matches) > 0:
             answer_text = matches[0].strip()
+
+            # if we don't expect a JSON output, return the answer text as is
+            if not json_output:
+                return answer_text
+
+            # otherwise, try to parse the answer text into a JSON object
             field_answers = self._check_bool_answer_text(answer_text)
             if field_answers is not None:
                 return field_answers
 
-        # finally, try taking all of the text; throw an exception if this doesn't work
-        field_answers = self._check_bool_answer_text(completion_text)
-        if field_answers is None:
-            raise Exception(f"Could not parse answer from completion text: {completion_text}")
+        # finally, try taking all of the text; for JSON output, throw an exception if parsing fails
+        if not json_output:
+            return completion_text
 
-        return field_answers
+        return self._check_bool_answer_text(completion_text, throw_exception=True)
 
     def _parse_answer(self, completion_text: str, fields: dict[str, FieldInfo] | None, json_output: bool, **kwargs) -> dict[str, list]:
         """Extract the answer from the completion object."""
@@ -275,8 +289,8 @@ class Generator(Generic[ContextType, InputType]):
 
         # extract the per-field answers from the completion text
         field_answers = (
-            self._parse_bool_answer(completion_text)
-            if self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()
+            self._parse_bool_answer(completion_text, json_output)
+            if self.prompt_strategy.is_filter_prompt() or self.prompt_strategy.is_join_prompt()
             else self._parse_convert_answer(completion_text, fields, json_output)
         )
 
@@ -299,6 +313,7 @@ class Generator(Generic[ContextType, InputType]):
 
         # generate a list of messages which can be used to construct a payload
         messages = self.prompt_factory.create_messages(candidate, fields, right_candidate, **kwargs)
+        is_audio_op = any(msg.get("type") == "input_audio" for msg in messages)
 
         # generate the text completion
         start_time = time.time()
@@ -307,7 +322,7 @@ class Generator(Generic[ContextType, InputType]):
             completion_kwargs = {}
             if not self.model.is_o_model() and not self.model.is_gpt_5_model():
                 completion_kwargs = {"temperature": kwargs.get("temperature", 0.0), **completion_kwargs}
-            if self.prompt_strategy.is_audio_prompt():
+            if is_audio_op:
                 completion_kwargs = {"modalities": ["text"], **completion_kwargs}
             if self.model.is_reasoning_model():
                 if self.model.is_vertex_model():
@@ -330,11 +345,10 @@ class Generator(Generic[ContextType, InputType]):
         # if there's an error generating the completion, we have to return an empty answer
         # and can only account for the time spent performing the failed generation
         except Exception as e:
-            print(f"Error generating completion: {e}")
             logger.error(f"Error generating completion: {e}")
             field_answers = (
                 {"passed_operator": False}
-                if self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()
+                if self.prompt_strategy.is_filter_prompt() or self.prompt_strategy.is_join_prompt()
                 else {field_name: None for field_name in fields}
             )
             reasoning = None
@@ -360,7 +374,7 @@ class Generator(Generic[ContextType, InputType]):
             #       for now, we only use tokens from prompt_token_details if it's an audio prompt
             # get output tokens (all text) and input tokens by modality
             output_tokens = usage["completion_tokens"]
-            if self.prompt_strategy.is_audio_prompt():
+            if is_audio_op:
                 input_audio_tokens = usage["prompt_tokens_details"].get("audio_tokens", 0)
                 input_text_tokens = usage["prompt_tokens_details"].get("text_tokens", 0)
                 input_image_tokens = 0
@@ -413,9 +427,9 @@ class Generator(Generic[ContextType, InputType]):
 
         # parse field answers
         field_answers = None 
-        if fields is not None and (self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()):
+        if fields is not None and (self.prompt_strategy.is_filter_prompt() or self.prompt_strategy.is_join_prompt()):
             field_answers = {"passed_operator": False}
-        elif fields is not None and not (self.prompt_strategy.is_bool_prompt() or self.prompt_strategy.is_join_prompt()):
+        elif fields is not None and not (self.prompt_strategy.is_filter_prompt() or self.prompt_strategy.is_join_prompt()):
             field_answers = {field_name: None for field_name in fields}
         try:
             field_answers = self._parse_answer(completion_text, fields, json_output, **kwargs)

palimpzest/query/operators/__init__.py

@@ -6,6 +6,8 @@ from palimpzest.query.operators.convert import ConvertOp as _ConvertOp
 from palimpzest.query.operators.convert import LLMConvert as _LLMConvert
 from palimpzest.query.operators.convert import LLMConvertBonded as _LLMConvertBonded
 from palimpzest.query.operators.convert import NonLLMConvert as _NonLLMConvert
+from palimpzest.query.operators.critique_and_refine import CritiqueAndRefineConvert as _CritiqueAndRefineConvert
+from palimpzest.query.operators.critique_and_refine import CritiqueAndRefineFilter as _CritiqueAndRefineFilter
 from palimpzest.query.operators.distinct import DistinctOp as _DistinctOp
 from palimpzest.query.operators.filter import FilterOp as _FilterOp
 from palimpzest.query.operators.filter import LLMFilter as _LLMFilter
@@ -46,12 +48,17 @@ from palimpzest.query.operators.logical import (
 from palimpzest.query.operators.logical import (
     RetrieveScan as _RetrieveScan,
 )
-from palimpzest.query.operators.mixture_of_agents_convert import MixtureOfAgentsConvert as _MixtureOfAgentsConvert
+from palimpzest.query.operators.mixture_of_agents import MixtureOfAgentsConvert as _MixtureOfAgentsConvert
+from palimpzest.query.operators.mixture_of_agents import MixtureOfAgentsFilter as _MixtureOfAgentsFilter
 from palimpzest.query.operators.physical import PhysicalOperator as _PhysicalOperator
 from palimpzest.query.operators.project import ProjectOp as _ProjectOp
+from palimpzest.query.operators.rag import RAGConvert as _RAGConvert
+from palimpzest.query.operators.rag import RAGFilter as _RAGFilter
 from palimpzest.query.operators.retrieve import RetrieveOp as _RetrieveOp
 from palimpzest.query.operators.scan import MarshalAndScanDataOp as _MarshalAndScanDataOp
 from palimpzest.query.operators.scan import ScanPhysicalOp as _ScanPhysicalOp
+from palimpzest.query.operators.split import SplitConvert as _SplitConvert
+from palimpzest.query.operators.split import SplitFilter as _SplitFilter
 
 LOGICAL_OPERATORS = [
     _LogicalOperator,
@@ -72,6 +79,8 @@ PHYSICAL_OPERATORS = (
     [_AggregateOp, _ApplyGroupByOp, _AverageAggregateOp, _CountAggregateOp]
     # convert
     + [_ConvertOp, _NonLLMConvert, _LLMConvert, _LLMConvertBonded]
+    # critique and refine
+    + [_CritiqueAndRefineConvert, _CritiqueAndRefineFilter]
     # distinct
     + [_DistinctOp]
     # scan
@@ -83,13 +92,17 @@ PHYSICAL_OPERATORS = (
     # limit
     + [_LimitScanOp]
     # mixture-of-agents
-    + [_MixtureOfAgentsConvert]
+    + [_MixtureOfAgentsConvert, _MixtureOfAgentsFilter]
     # physical
     + [_PhysicalOperator]
     # project
     + [_ProjectOp]
+    # rag
+    + [_RAGConvert, _RAGFilter]
     # retrieve
     + [_RetrieveOp]
+    # split
+    + [_SplitConvert, _SplitFilter]
 )
 
 __all__ = [

palimpzest/query/operators/aggregate.py

@@ -113,18 +113,20 @@ class ApplyGroupByOp(AggregateOp):
         group_by_fields = self.group_by_sig.group_by_fields
         agg_fields = self.group_by_sig.get_agg_field_names()
         for g in agg_state:
-            dr = DataRecord.from_agg_parents(
-                schema=self.group_by_sig.output_schema(),
-                parent_records=candidates,
-            )
+            # build up data item
+            data_item = {}
             for i in range(0, len(g)):
                 k = g[i]
-                setattr(dr, group_by_fields[i], k)
+                data_item[group_by_fields[i]] = k
             vals = agg_state[g]
             for i in range(0, len(vals)):
                 v = ApplyGroupByOp.agg_final(self.group_by_sig.agg_funcs[i], vals[i])
-                setattr(dr, agg_fields[i], v)
+                data_item[agg_fields[i]] = v
 
+            # create new DataRecord
+            schema = self.group_by_sig.output_schema()
+            data_item = schema(**data_item)
+            dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)
             drs.append(dr)
 
         # create RecordOpStats objects
@@ -132,9 +134,9 @@ class ApplyGroupByOp(AggregateOp):
         record_op_stats_lst = []
         for dr in drs:
             record_op_stats = RecordOpStats(
-                record_id=dr.id,
-                record_parent_ids=dr.parent_ids,
-                record_source_indices=dr.source_indices,
+                record_id=dr._id,
+                record_parent_ids=dr._parent_ids,
+                record_source_indices=dr._source_indices,
                 record_state=dr.to_dict(include_bytes=False),
                 full_op_id=self.get_full_op_id(),
                 logical_op_id=self.logical_op_id,
@@ -197,7 +199,6 @@ class AverageAggregateOp(AggregateOp):
         # NOTE: right now we perform a check in the constructor which enforces that the input_schema
         #       has a single field which is numeric in nature; in the future we may want to have a
         #       cleaner way of computing the value (rather than `float(list(candidate...))` below)
-        dr = DataRecord.from_agg_parents(schema=Average, parent_records=candidates)
         summation, total = 0, 0
         for candidate in candidates:
             try:
@@ -205,13 +206,14 @@ class AverageAggregateOp(AggregateOp):
                 total += 1
             except Exception:
                 pass
-        dr.average = summation / total
+        data_item = Average(average=summation / total)
+        dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
-            record_id=dr.id,
-            record_parent_ids=dr.parent_ids,
-            record_source_indices=dr.source_indices,
+            record_id=dr._id,
+            record_parent_ids=dr._parent_ids,
+            record_source_indices=dr._source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,
@@ -260,14 +262,14 @@ class CountAggregateOp(AggregateOp):
         start_time = time.time()
 
         # create new DataRecord
-        dr = DataRecord.from_agg_parents(schema=Count, parent_records=candidates)
-        dr.count = len(candidates)
+        data_item = Count(count=len(candidates))
+        dr = DataRecord.from_agg_parents(data_item, parent_records=candidates)
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
-            record_id=dr.id,
-            record_parent_ids=dr.parent_ids,
-            record_source_indices=dr.source_indices,
+            record_id=dr._id,
+            record_parent_ids=dr._parent_ids,
+            record_source_indices=dr._source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,

palimpzest/query/operators/compute.py

@@ -93,17 +93,15 @@ class SmolAgentsCompute(PhysicalOperator):
         Given an input DataRecord and a determination of whether it passed the filter or not,
         construct the resulting RecordSet.
         """
-        # create new DataRecord and set passed_operator attribute
-        dr = DataRecord.from_parent(self.output_schema, parent_record=candidate)
-        for field in self.output_schema.model_fields:
-            if field in answer:
-                dr[field] = answer[field]
+        # create new DataRecord
+        data_item = {field: answer[field] for field in self.output_schema.model_fields if field in answer}
+        dr = DataRecord.from_parent(self.output_schema, data_item, parent_record=candidate)
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
-            record_id=dr.id,
-            record_parent_ids=dr.parent_ids,
-            record_source_indices=dr.source_indices,
+            record_id=dr._id,
+            record_parent_ids=dr._parent_ids,
+            record_source_indices=dr._source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,

palimpzest/query/operators/convert.py

@@ -74,25 +74,14 @@ class ConvertOp(PhysicalOperator, ABC):
 
         drs = []
         for idx in range(max(n_records, 1)):
-            # initialize record with the correct output schema, parent record, and cardinality idx
-            dr = DataRecord.from_parent(self.output_schema, parent_record=candidate, cardinality_idx=idx)
-
-            # copy all fields from the input record
-            # NOTE: this means that records processed by PZ converts will inherit all pre-computed fields
-            #       in an incremental fashion; this is a design choice which may be revisited in the future
-            for field in candidate.get_field_names():
-                setattr(dr, field, getattr(candidate, field))
-
-            # get input field names and output field names
-            input_fields = list(self.input_schema.model_fields)
-            output_fields = list(self.output_schema.model_fields)
-
             # parse newly generated fields from the field_answers dictionary for this field; if the list
             # of generated values is shorter than the number of records, we fill in with None
-            for field in output_fields:
-                if field not in input_fields:
-                    value = field_answers[field][idx] if idx < len(field_answers[field]) else None
-                    setattr(dr, field, value)
+            data_item = {}
+            for field in self.generated_fields:
+                data_item[field] = field_answers[field][idx] if idx < len(field_answers[field]) else None
+
+            # initialize record with the correct output schema, data_item, parent record, and cardinality idx
+            dr = DataRecord.from_parent(self.output_schema, data_item, parent_record=candidate, cardinality_idx=idx)
 
             # append data record to list of output data records
             drs.append(dr)
@@ -117,9 +106,9 @@ class ConvertOp(PhysicalOperator, ABC):
         # create the RecordOpStats objects for each output record
         record_op_stats_lst = [
             RecordOpStats(
-                record_id=dr.id,
-                record_parent_ids=dr.parent_ids,
-                record_source_indices=dr.source_indices,
+                record_id=dr._id,
+                record_parent_ids=dr._parent_ids,
+                record_source_indices=dr._source_indices,
                 record_state=dr.to_dict(include_bytes=False),
                 full_op_id=self.get_full_op_id(),
                 logical_op_id=self.logical_op_id,
@@ -127,7 +116,7 @@ class ConvertOp(PhysicalOperator, ABC):
                 time_per_record=time_per_record,
                 cost_per_record=per_record_stats.cost_per_record,
                 model_name=self.get_model_name(),
-                answer={field_name: getattr(dr, field_name) for field_name in field_names},
+                answer={field_name: getattr(dr, field_name, None) for field_name in field_names},
                 input_fields=list(self.input_schema.model_fields),
                 generated_fields=field_names,
                 total_input_tokens=per_record_stats.total_input_tokens,
@@ -139,7 +128,6 @@ class ConvertOp(PhysicalOperator, ABC):
                 total_llm_calls=per_record_stats.total_llm_calls,
                 total_embedding_llm_calls=per_record_stats.total_embedding_llm_calls,
                 failed_convert=(not successful_convert),
-                image_operation=self.is_image_conversion(),
                 op_details={k: str(v) for k, v in self.get_id_params().items()},
             )
             for dr in records
@@ -148,11 +136,6 @@ class ConvertOp(PhysicalOperator, ABC):
         # create and return the DataRecordSet
         return DataRecordSet(records, record_op_stats_lst)
 
-    @abstractmethod
-    def is_image_conversion(self) -> bool:
-        """Return True if the convert operation processes an image, False otherwise."""
-        pass
-
     @abstractmethod
     def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[str, list], GenerationStats]:
         """
@@ -216,11 +199,6 @@ class NonLLMConvert(ConvertOp):
         op += f"    UDF: {self.udf.__name__}\n"
         return op
 
-    def is_image_conversion(self) -> bool:
-        # NOTE: even if the UDF is processing an image, we do not consider this an image conversion
-        # (the output of this function will be used by the CostModel in a way which does not apply to UDFs)
-        return False
-
     def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
         """
         Compute naive cost estimates for the NonLLMConvert operation. These estimates assume
@@ -287,7 +265,7 @@ class LLMConvert(ConvertOp):
     def __init__(
         self,
         model: Model,
-        prompt_strategy: PromptStrategy = PromptStrategy.COT_QA,
+        prompt_strategy: PromptStrategy = PromptStrategy.MAP,
         reasoning_effort: str | None = None,
         *args,
         **kwargs,
@@ -330,9 +308,6 @@ class LLMConvert(ConvertOp):
     def get_model_name(self):
         return None if self.model is None else self.model.value
 
-    def is_image_conversion(self) -> bool:
-        return self.prompt_strategy.is_image_prompt()
-
     def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
         """
         Compute naive cost estimates for the LLMConvert operation. Implicitly, these estimates
@@ -350,7 +325,7 @@ class LLMConvert(ConvertOp):
 
         # get est. of conversion cost (in USD) per record from model card
         usd_per_input_token = MODEL_CARDS[model_name].get("usd_per_input_token")
-        if getattr(self, "prompt_strategy", None) is not None and self.prompt_strategy.is_audio_prompt():
+        if getattr(self, "prompt_strategy", None) is not None and self.is_audio_op():
             usd_per_input_token = MODEL_CARDS[model_name]["usd_per_audio_input_token"]
 
         model_conversion_usd_per_record = (

palimpzest/query/operators/critique_and_refine.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic.fields import FieldInfo
+
+from palimpzest.constants import MODEL_CARDS, Cardinality, Model, PromptStrategy
+from palimpzest.core.elements.records import DataRecord
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates
+from palimpzest.query.generators.generators import Generator
+from palimpzest.query.operators.convert import LLMConvert
+from palimpzest.query.operators.filter import LLMFilter
+
+# TYPE DEFINITIONS
+FieldName = str
+
+
+class CritiqueAndRefineConvert(LLMConvert):
+
+    def __init__(
+        self,
+        critic_model: Model,
+        refine_model: Model,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.critic_model = critic_model
+        self.refine_model = refine_model
+
+        # create generators
+        self.critic_generator = Generator(self.critic_model, PromptStrategy.MAP_CRITIC, self.reasoning_effort, self.api_base, self.cardinality, self.desc, self.verbose)
+        self.refine_generator = Generator(self.refine_model, PromptStrategy.MAP_REFINE, self.reasoning_effort, self.api_base, self.cardinality, self.desc, self.verbose)
+
+    def __str__(self):
+        op = super().__str__()
+        op += f"    Critic Model: {self.critic_model}\n"
+        op += f"    Refine Model: {self.refine_model}\n"
+        return op
+
+    def get_id_params(self):
+        id_params = super().get_id_params()
+        id_params = {
+            "critic_model": self.critic_model.value,
+            "refine_model": self.refine_model.value,
+            **id_params,
+        }
+
+        return id_params
+
+    def get_op_params(self):
+        op_params = super().get_op_params()
+        op_params = {
+            "critic_model": self.critic_model,
+            "refine_model": self.refine_model,
+            **op_params,
+        }
+
+        return op_params
+
+    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
+        """
+        Currently, we are invoking `self.model`, then critiquing its output with `self.critic_model`, and
+        finally refining the output with `self.refine_model`. Thus, we roughly expect to incur the cost
+        and time of three LLMConverts. In practice, this naive quality estimate will be overwritten by the
+        CostModel's estimate once it executes a few instances of the operator.
+        """
+        # get naive cost estimates for first LLM call and multiply by 3 for now;
+        # of course we should sum individual estimates for each model, but this is a rough estimate
+        # and in practice we will need to revamp our naive cost estimates in the near future
+        naive_op_cost_estimates = 3 * super().naive_cost_estimates(source_op_cost_estimates)
+
+        # for naive setting, estimate quality as quality of refine model
+        model_quality = MODEL_CARDS[self.refine_model.value]["overall"] / 100.0
+        naive_op_cost_estimates.quality = model_quality
+        naive_op_cost_estimates.quality_lower_bound = naive_op_cost_estimates.quality
+        naive_op_cost_estimates.quality_upper_bound = naive_op_cost_estimates.quality
+
+        return naive_op_cost_estimates
+
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[FieldName, list[Any]], GenerationStats]:
+        # get input fields
+        input_fields = self.get_input_fields()
+
+        # NOTE: when I merge in the `abacus` branch, I will want to update this to reflect the changes I made to reasoning extraction
+        # execute the initial model
+        original_gen_kwargs = {"project_cols": input_fields, "output_schema": self.output_schema}
+        field_answers, reasoning, original_gen_stats, original_messages = self.generator(candidate, fields, **original_gen_kwargs)
+        original_output = f"REASONING: {reasoning}\nANSWER: {field_answers}\n"
+
+        # execute the critic model
+        critic_gen_kwargs = {"original_output": original_output, "original_messages": original_messages, **original_gen_kwargs}
+        _, reasoning, critic_gen_stats, _ = self.critic_generator(candidate, fields, json_output=False, **critic_gen_kwargs)
+        critique_output = f"CRITIQUE: {reasoning}\n"
+
+        # execute the refinement model
+        refine_gen_kwargs = {"critique_output": critique_output, **critic_gen_kwargs}
+        field_answers, reasoning, refine_gen_stats, _ = self.refine_generator(candidate, fields, **refine_gen_kwargs)
+
+        # compute the total generation stats
+        generation_stats = original_gen_stats + critic_gen_stats + refine_gen_stats
+
+        return field_answers, generation_stats
+
+
+class CritiqueAndRefineFilter(LLMFilter):
+
+    def __init__(
+        self,
+        critic_model: Model,
+        refine_model: Model,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.critic_model = critic_model
+        self.refine_model = refine_model
+
+        # create generators
+        self.critic_generator = Generator(self.critic_model, PromptStrategy.FILTER_CRITIC, self.reasoning_effort, self.api_base, Cardinality.ONE_TO_ONE, self.desc, self.verbose)
+        self.refine_generator = Generator(self.refine_model, PromptStrategy.FILTER_REFINE, self.reasoning_effort, self.api_base, Cardinality.ONE_TO_ONE, self.desc, self.verbose)
+
+    def __str__(self):
+        op = super().__str__()
+        op += f"    Critic Model: {self.critic_model}\n"
+        op += f"    Refine Model: {self.refine_model}\n"
+        return op
+
+    def get_id_params(self):
+        id_params = super().get_id_params()
+        id_params = {
+            "critic_model": self.critic_model.value,
+            "refine_model": self.refine_model.value,
+            **id_params,
+        }
+
+        return id_params
+
+    def get_op_params(self):
+        op_params = super().get_op_params()
+        op_params = {
+            "critic_model": self.critic_model,
+            "refine_model": self.refine_model,
+            **op_params,
+        }
+
+        return op_params
+
+    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
+        """
+        Currently, we are invoking `self.model`, then critiquing its output with `self.critic_model`, and
+        finally refining the output with `self.refine_model`. Thus, we roughly expect to incur the cost
+        and time of three LLMFilters. In practice, this naive quality estimate will be overwritten by the
+        CostModel's estimate once it executes a few instances of the operator.
+        """
+        # get naive cost estimates for first LLM call and multiply by 3 for now;
+        # of course we should sum individual estimates for each model, but this is a rough estimate
+        # and in practice we will need to revamp our naive cost estimates in the near future
+        naive_op_cost_estimates = 3 * super().naive_cost_estimates(source_op_cost_estimates)
+
+        # for naive setting, estimate quality as quality of refine model
+        model_quality = MODEL_CARDS[self.refine_model.value]["overall"] / 100.0
+        naive_op_cost_estimates.quality = model_quality
+        naive_op_cost_estimates.quality_lower_bound = naive_op_cost_estimates.quality
+        naive_op_cost_estimates.quality_upper_bound = naive_op_cost_estimates.quality
+
+        return naive_op_cost_estimates
+
+    def filter(self, candidate: DataRecord) -> tuple[dict[str, bool], GenerationStats]:
+        # get input fields
+        input_fields = self.get_input_fields()
+
+        # construct output fields
+        fields = {"passed_operator": FieldInfo(annotation=bool, description="Whether the record passed the filter operation")}
+
+        # NOTE: when I merge in the `abacus` branch, I will want to update this to reflect the changes I made to reasoning extraction
+        # execute the initial model
+        original_gen_kwargs = {"project_cols": input_fields, "filter_condition": self.filter_obj.filter_condition}
+        field_answers, reasoning, original_gen_stats, original_messages = self.generator(candidate, fields, **original_gen_kwargs)
+        original_output = f"REASONING: {reasoning}\nANSWER: {str(field_answers['passed_operator']).upper()}\n"
+
+        # execute the critic model
+        critic_gen_kwargs = {"original_output": original_output, "original_messages": original_messages, **original_gen_kwargs}
+        _, reasoning, critic_gen_stats, _ = self.critic_generator(candidate, fields, json_output=False, **critic_gen_kwargs)
+        critique_output = f"CRITIQUE: {reasoning}\n"
+
+        # execute the refinement model
+        refine_gen_kwargs = {"critique_output": critique_output, **critic_gen_kwargs}
+        field_answers, reasoning, refine_gen_stats, _ = self.refine_generator(candidate, fields, **refine_gen_kwargs)
+
+        # compute the total generation stats
+        generation_stats = original_gen_stats + critic_gen_stats + refine_gen_stats
+
+        return field_answers, generation_stats

palimpzest/query/operators/distinct.py

@@ -35,27 +35,27 @@ class DistinctOp(PhysicalOperator):
 
     def __call__(self, candidate: DataRecord) -> DataRecordSet:
         # create new DataRecord
-        dr = DataRecord.from_parent(schema=candidate.schema, parent_record=candidate)
+        dr = DataRecord.from_parent(schema=candidate.schema, data_item={}, parent_record=candidate)
 
         # output record only if it has not been seen before
         record_str = dr.to_json_str(project_cols=self.distinct_cols, bytes_to_str=True, sorted=True)
         record_hash = f"{hash(record_str)}"
-        dr.passed_operator = record_hash not in self._distinct_seen
-        if dr.passed_operator:
+        dr._passed_operator = record_hash not in self._distinct_seen
+        if dr._passed_operator:
             self._distinct_seen.add(record_hash)
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
-            record_id=dr.id,
-            record_parent_ids=dr.parent_ids,
-            record_source_indices=dr.source_indices,
+            record_id=dr._id,
+            record_parent_ids=dr._parent_ids,
+            record_source_indices=dr._source_indices,
             record_state=dr.to_dict(include_bytes=False),
             full_op_id=self.get_full_op_id(),
             logical_op_id=self.logical_op_id,
             op_name=self.op_name(),
             time_per_record=0.0,
             cost_per_record=0.0,
-            passed_operator=dr.passed_operator,
+            passed_operator=dr._passed_operator,
             op_details={k: str(v) for k, v in self.get_id_params().items()},
         )