palimpzest 0.7.20__py3-none-any.whl → 0.8.0__py3-none-any.whl

palimpzest/query/operators/__init__.py

@@ -6,9 +6,12 @@ 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.distinct import DistinctOp as _DistinctOp
 from palimpzest.query.operators.filter import FilterOp as _FilterOp
 from palimpzest.query.operators.filter import LLMFilter as _LLMFilter
 from palimpzest.query.operators.filter import NonLLMFilter as _NonLLMFilter
+from palimpzest.query.operators.join import JoinOp as _JoinOp
+from palimpzest.query.operators.join import NestedLoopsJoin as _NestedLoopsJoin
 from palimpzest.query.operators.limit import LimitScanOp as _LimitScanOp
 from palimpzest.query.operators.logical import (
     Aggregate as _Aggregate,
@@ -17,10 +20,10 @@ from palimpzest.query.operators.logical import (
     BaseScan as _BaseScan,
 )
 from palimpzest.query.operators.logical import (
-    CacheScan as _CacheScan,
+    ConvertScan as _ConvertScan,
 )
 from palimpzest.query.operators.logical import (
-    ConvertScan as _ConvertScan,
+    Distinct as _Distinct,
 )
 from palimpzest.query.operators.logical import (
     FilteredScan as _FilteredScan,
@@ -28,6 +31,9 @@ from palimpzest.query.operators.logical import (
 from palimpzest.query.operators.logical import (
     GroupByAggregate as _GroupByAggregate,
 )
+from palimpzest.query.operators.logical import (
+    JoinOp as _LogicalJoinOp,
+)
 from palimpzest.query.operators.logical import (
     LimitScan as _LimitScan,
 )
@@ -44,7 +50,6 @@ from palimpzest.query.operators.mixture_of_agents_convert import MixtureOfAgents
 from palimpzest.query.operators.physical import PhysicalOperator as _PhysicalOperator
 from palimpzest.query.operators.project import ProjectOp as _ProjectOp
 from palimpzest.query.operators.retrieve import RetrieveOp as _RetrieveOp
-from palimpzest.query.operators.scan import CacheScanDataOp as _CacheScanDataOp
 from palimpzest.query.operators.scan import MarshalAndScanDataOp as _MarshalAndScanDataOp
 from palimpzest.query.operators.scan import ScanPhysicalOp as _ScanPhysicalOp
 
@@ -52,10 +57,11 @@ LOGICAL_OPERATORS = [
     _LogicalOperator,
     _Aggregate,
     _BaseScan,
-    _CacheScan,
     _ConvertScan,
+    _Distinct,
     _FilteredScan,
     _GroupByAggregate,
+    _LogicalJoinOp,
     _LimitScan,
     _Project,
     _RetrieveScan,
@@ -66,10 +72,14 @@ PHYSICAL_OPERATORS = (
     [_AggregateOp, _ApplyGroupByOp, _AverageAggregateOp, _CountAggregateOp]
     # convert
     + [_ConvertOp, _NonLLMConvert, _LLMConvert, _LLMConvertBonded]
+    # distinct
+    + [_DistinctOp]
     # scan
-    + [_ScanPhysicalOp, _MarshalAndScanDataOp, _CacheScanDataOp]
+    + [_ScanPhysicalOp, _MarshalAndScanDataOp]
     # filter
     + [_FilterOp, _NonLLMFilter, _LLMFilter]
+    # join
+    + [_JoinOp, _NestedLoopsJoin]
     # limit
     + [_LimitScanOp]
     # mixture-of-agents

palimpzest/query/operators/aggregate.py

@@ -3,10 +3,10 @@ from __future__ import annotations
 import time
 
 from palimpzest.constants import NAIVE_EST_NUM_GROUPS, AggFunc
-from palimpzest.core.data.dataclasses import OperatorCostEstimates, RecordOpStats
 from palimpzest.core.elements.groupbysig import GroupBySig
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
-from palimpzest.core.lib.schemas import Number
+from palimpzest.core.lib.schemas import Average, Count
+from palimpzest.core.models import OperatorCostEstimates, RecordOpStats
 from palimpzest.query.operators.physical import PhysicalOperator
 
 
@@ -16,7 +16,7 @@ class AggregateOp(PhysicalOperator):
     __call__ methods. Thus, we use a slightly modified abstract base class for
     these operators.
     """
-    def __call__(self, candidates: DataRecordSet) -> DataRecordSet:
+    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
         raise NotImplementedError("Using __call__ from abstract method")
 
 
@@ -67,6 +67,8 @@ class ApplyGroupByOp(AggregateOp):
             return state + 1
         elif func.lower() == "average":
             sum, cnt = state
+            if val is None:
+                return (sum, cnt)
             return (sum + val, cnt + 1)
         else:
             raise Exception("Unknown agg function " + func)
@@ -77,11 +79,11 @@ class ApplyGroupByOp(AggregateOp):
             return state
         elif func.lower() == "average":
             sum, cnt = state
-            return float(sum) / cnt
+            return float(sum) / cnt if cnt > 0 else None
         else:
             raise Exception("Unknown agg function " + func)
 
-    def __call__(self, candidates: DataRecordSet) -> DataRecordSet:
+    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
         start_time = time.time()
 
         # build group array
@@ -107,17 +109,13 @@ class ApplyGroupByOp(AggregateOp):
             agg_state[group] = state
 
         # return list of data records (one per group)
-        drs = []
+        drs: list[DataRecord] = []
         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(self.group_by_sig.output_schema())
-            # NOTE: this will set the parent_id and source_idx to be the id of the final source record;
-            #       in the near future we may want to have parent_id accept a list of ids
-            dr = DataRecord.from_parent(
+            dr = DataRecord.from_agg_parents(
                 schema=self.group_by_sig.output_schema(),
-                parent_record=candidates[-1],
-                project_cols=[],
+                parent_records=candidates,
             )
             for i in range(0, len(g)):
                 k = g[i]
@@ -135,8 +133,8 @@ class ApplyGroupByOp(AggregateOp):
         for dr in drs:
             record_op_stats = RecordOpStats(
                 record_id=dr.id,
-                record_parent_id=dr.parent_id,
-                record_source_idx=dr.source_idx,
+                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,
@@ -155,13 +153,20 @@ class AverageAggregateOp(AggregateOp):
     # NOTE: we don't actually need / use agg_func here (yet)
 
     def __init__(self, agg_func: AggFunc, *args, **kwargs):
-        kwargs["output_schema"] = Number
+        # enforce that output schema is correct
+        assert kwargs["output_schema"] == Average, "AverageAggregateOp requires output_schema to be Average"
+
+        # enforce that input schema is a single numeric field
+        input_field_types = list(kwargs["input_schema"].model_fields.values())
+        assert len(input_field_types) == 1, "AverageAggregateOp requires input_schema to have exactly one field"
+        numeric_field_types = [bool, int, float, bool | None, int | None, float | None, int | float, int | float | None]
+        is_numeric = input_field_types[0].annotation in numeric_field_types
+        assert is_numeric, f"AverageAggregateOp requires input_schema to have a numeric field type, i.e. one of: {numeric_field_types}\nGot: {input_field_types[0]}"
+
+        # call parent constructor
         super().__init__(*args, **kwargs)
         self.agg_func = agg_func
 
-        if not self.input_schema.get_desc() == Number.get_desc():
-            raise Exception("Aggregate function AVERAGE is only defined over Numbers")
-
     def __str__(self):
         op = super().__str__()
         op += f"    Function: {str(self.agg_func)}\n"
@@ -184,19 +189,29 @@ class AverageAggregateOp(AggregateOp):
             quality=1.0,
         )
 
-    def __call__(self, candidates: DataRecordSet) -> DataRecordSet:
+    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
         start_time = time.time()
 
-        # NOTE: this will set the parent_id and source_idx to be the id of the final source record;
-        #       in the near future we may want to have parent_id accept a list of ids
-        dr = DataRecord.from_parent(schema=Number, parent_record=candidates[-1], project_cols=[])
-        dr.value = sum(list(map(lambda c: float(c.value), candidates))) / len(candidates)
+        # NOTE: we currently do not guarantee that input values conform to their specified type;
+        #       as a result, we simply omit any values which do not parse to a float from the average
+        # 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:
+                summation += float(list(candidate.to_dict().values())[0])
+                total += 1
+            except Exception:
+                pass
+        dr.average = summation / total
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
             record_id=dr.id,
-            record_parent_id=dr.parent_id,
-            record_source_idx=dr.source_idx,
+            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,
@@ -212,7 +227,10 @@ class CountAggregateOp(AggregateOp):
     # NOTE: we don't actually need / use agg_func here (yet)
 
     def __init__(self, agg_func: AggFunc, *args, **kwargs):
-        kwargs["output_schema"] = Number
+        # enforce that output schema is correct
+        assert kwargs["output_schema"] == Count, "CountAggregateOp requires output_schema to be Count"
+
+        # call parent constructor
         super().__init__(*args, **kwargs)
         self.agg_func = agg_func
 
@@ -238,19 +256,18 @@ class CountAggregateOp(AggregateOp):
             quality=1.0,
         )
 
-    def __call__(self, candidates: DataRecordSet) -> DataRecordSet:
+    def __call__(self, candidates: list[DataRecord]) -> DataRecordSet:
         start_time = time.time()
 
-        # NOTE: this will set the parent_id to be the id of the final source record;
-        #       in the near future we may want to have parent_id accept a list of ids
-        dr = DataRecord.from_parent(schema=Number, parent_record=candidates[-1], project_cols=[])
-        dr.value = len(candidates)
+        # create new DataRecord
+        dr = DataRecord.from_agg_parents(schema=Count, parent_records=candidates)
+        dr.count = len(candidates)
 
         # create RecordOpStats object
         record_op_stats = RecordOpStats(
             record_id=dr.id,
-            record_parent_id=dr.parent_id,
-            record_source_idx=dr.source_idx,
+            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

@@ -0,0 +1,201 @@
+import functools
+import inspect
+import os
+import time
+from typing import Any
+
+from smolagents import CodeAgent, LiteLLMModel, tool
+
+from palimpzest.core.data.context import Context
+from palimpzest.core.data.context_manager import ContextManager
+from palimpzest.core.elements.records import DataRecord, DataRecordSet
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates, RecordOpStats
+from palimpzest.query.operators.physical import PhysicalOperator
+
+# TODO: need to store final executed code in compute() operator so that humans can debug when human-in-the-loop
+
+def make_tool(bound_method):
+    # Get the original function and bound instance
+    func = bound_method.__func__
+    instance = bound_method.__self__
+    
+    # Get the signature and remove 'self'
+    sig = inspect.signature(func)
+    params = list(sig.parameters.values())[1:]  # skip 'self'
+    new_sig = inspect.Signature(parameters=params, return_annotation=sig.return_annotation)
+
+    # Create a wrapper function dynamically
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(instance, *args, **kwargs)
+
+    # Update the __signature__ to reflect the new one without 'self'
+    wrapper.__signature__ = new_sig
+
+    return wrapper
+
+
+class SmolAgentsCompute(PhysicalOperator):
+    """
+    """
+    def __init__(self, context_id: str, instruction: str, additional_contexts: list[Context] | None = None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.context_id = context_id
+        self.instruction = instruction
+        self.additional_contexts = [] if additional_contexts is None else additional_contexts
+        # self.model_id = "anthropic/claude-3-7-sonnet-latest"
+        self.model_id = "openai/gpt-4o-mini-2024-07-18"
+        # self.model_id = "openai/gpt-4o-2024-08-06"
+        api_key = os.getenv("ANTHROPIC_API_KEY") if "anthropic" in self.model_id else os.getenv("OPENAI_API_KEY")
+        self.model = LiteLLMModel(model_id=self.model_id, api_key=api_key)
+
+    def __str__(self):
+        op = super().__str__()
+        op += f"    Context ID: {self.context_id:20s}\n"
+        op += f"    Instruction: {self.instruction:20s}\n"
+        op += f"    Add. Ctxs: {self.additional_contexts}\n"
+        return op
+
+    def get_id_params(self):
+        id_params = super().get_id_params()
+        return {
+            "context_id": self.context_id,
+            "instruction": self.instruction,
+            "additional_contexts": self.additional_contexts,
+            **id_params,
+        }
+
+    def get_op_params(self):
+        op_params = super().get_op_params()
+        return {
+            "context_id": self.context_id,
+            "instruction": self.instruction,
+            "additional_contexts": self.additional_contexts,
+            **op_params,
+        }
+
+    def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
+        return OperatorCostEstimates(
+            cardinality=source_op_cost_estimates.cardinality,
+            time_per_record=100,
+            cost_per_record=1,
+            quality=1.0,
+        )
+
+    def _create_record_set(
+        self,
+        candidate: DataRecord,
+        generation_stats: GenerationStats,
+        total_time: float,
+        answer: dict[str, Any],
+    ) -> DataRecordSet:
+        """
+        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 RecordOpStats object
+        record_op_stats = RecordOpStats(
+            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=total_time,
+            cost_per_record=generation_stats.cost_per_record,
+            model_name=self.get_model_name(),
+            total_input_tokens=generation_stats.total_input_tokens,
+            total_output_tokens=generation_stats.total_output_tokens,
+            total_input_cost=generation_stats.total_input_cost,
+            total_output_cost=generation_stats.total_output_cost,
+            llm_call_duration_secs=generation_stats.llm_call_duration_secs,
+            fn_call_duration_secs=generation_stats.fn_call_duration_secs,
+            total_llm_calls=generation_stats.total_llm_calls,
+            total_embedding_llm_calls=generation_stats.total_embedding_llm_calls,
+            answer={k: v.description if isinstance(v, Context) else v for k, v in answer.items()},
+            op_details={k: str(v) for k, v in self.get_id_params().items()},
+        )
+
+        return DataRecordSet([dr], [record_op_stats])
+
+    def __call__(self, candidate: DataRecord) -> Any:
+        start_time = time.time()
+
+        # get the input context object and its tools
+        input_context: Context = candidate.context
+        description = input_context.description
+        tools = [tool(make_tool(f)) for f in input_context.tools]
+
+        # update the description to include any additional contexts
+        for ctx in self.additional_contexts:
+            # TODO: remove additional context if it is an ancestor of the input context
+            # (not just if it is equal to the input context)
+            if ctx.id == input_context.id:
+                continue
+            description += f"\n\nHere is some additional Context which may be useful:\n\n{ctx.description}"
+
+        # perform the computation
+        instructions = f"\n\nHere is a description of the Context whose data you will be working with, as well as any previously computed results:\n\n{description}"
+        agent = CodeAgent(
+            tools=tools,
+            model=self.model,
+            add_base_tools=False,
+            instructions=instructions,
+            return_full_result=True,
+            additional_authorized_imports=["pandas", "io", "os"],
+            planning_interval=4,
+            max_steps=30,
+        )
+        result = agent.run(self.instruction)
+        # NOTE: you can see the system prompt with `agent.memory.system_prompt.system_prompt`
+        # full_steps = agent.memory.get_full_steps()
+
+        # compute generation stats
+        response = result.output
+        input_tokens = result.token_usage.input_tokens
+        output_tokens = result.token_usage.output_tokens
+        cost_per_input_token = (3.0 / 1e6) if "anthropic" in self.model_id else (0.15 / 1e6) # (2.5 / 1e6) #
+        cost_per_output_token = (15.0 / 1e6) if "anthropic" in self.model_id else (0.6 / 1e6) # (10.0 / 1e6) #
+        input_cost = input_tokens * cost_per_input_token
+        output_cost = output_tokens * cost_per_output_token
+        generation_stats = GenerationStats(
+            model_name=self.model_id,
+            total_input_tokens=input_tokens,
+            total_output_tokens=output_tokens,
+            total_input_cost=input_cost,
+            total_output_cost=output_cost,
+            cost_per_record=input_cost + output_cost,
+            llm_call_duration_secs=time.time() - start_time,
+        )
+
+        # update the description of the computed Context to include the result
+        new_description = f"RESULT: {response}\n\n"
+        cm = ContextManager()
+        cm.update_context(id=self.context_id, description=new_description)
+
+        # create and return record set
+        field_answers = {
+            "context": cm.get_context(id=self.context_id),
+            f"result-{self.context_id}": response,
+        }
+        record_set = self._create_record_set(
+            candidate,
+            generation_stats,
+            time.time() - start_time,
+            field_answers,
+        )
+
+        return record_set
+
+# import json; json.dumps(agent.memory.get_full_steps())
+# agent.memory.get_full_steps()[1].keys()
+# dict_keys(['step_number', 'timing', 'model_input_messages', 'tool_calls', 'error', 'model_output_message', 'model_output', 'code_action', 'observations', 'observations_images', 
+# 'action_output', 'token_usage', 'is_final_answer'])
+# agent.memory.get_full_steps()[1]['action_output']

palimpzest/query/operators/convert.py

@@ -4,6 +4,8 @@ import time
 from abc import ABC, abstractmethod
 from typing import Callable
 
+from pydantic.fields import FieldInfo
+
 from palimpzest.constants import (
     MODEL_CARDS,
     NAIVE_EST_NUM_INPUT_TOKENS,
@@ -13,12 +15,10 @@ from palimpzest.constants import (
     Model,
     PromptStrategy,
 )
-from palimpzest.core.data.dataclasses import GenerationStats, OperatorCostEstimates, RecordOpStats
 from palimpzest.core.elements.records import DataRecord, DataRecordSet
-from palimpzest.core.lib.fields import Field
-from palimpzest.query.generators.generators import generator_factory
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates, RecordOpStats
+from palimpzest.query.generators.generators import Generator
 from palimpzest.query.operators.physical import PhysicalOperator
-from palimpzest.utils.model_helpers import get_vision_models
 
 
 class ConvertOp(PhysicalOperator, ABC):
@@ -26,14 +26,12 @@ class ConvertOp(PhysicalOperator, ABC):
         self,
         cardinality: Cardinality = Cardinality.ONE_TO_ONE,
         udf: Callable | None = None,
-        desc: str | None = None,
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.cardinality = cardinality
         self.udf = udf
-        self.desc = desc
 
     def get_id_params(self):
         id_params = super().get_id_params()
@@ -47,7 +45,7 @@ class ConvertOp(PhysicalOperator, ABC):
 
     def get_op_params(self):
         op_params = super().get_op_params()
-        op_params = {"cardinality": self.cardinality, "udf": self.udf, "desc": self.desc, **op_params}
+        op_params = {"cardinality": self.cardinality, "udf": self.udf, **op_params}
 
         return op_params
 
@@ -78,8 +76,8 @@ class ConvertOp(PhysicalOperator, ABC):
                 setattr(dr, field, getattr(candidate, field))
 
             # get input field names and output field names
-            input_fields = self.input_schema.field_names()
-            output_fields = self.output_schema.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
@@ -112,8 +110,8 @@ class ConvertOp(PhysicalOperator, ABC):
         record_op_stats_lst = [
             RecordOpStats(
                 record_id=dr.id,
-                record_parent_id=dr.parent_id,
-                record_source_idx=dr.source_idx,
+                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,
@@ -122,7 +120,7 @@ class ConvertOp(PhysicalOperator, ABC):
                 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},
-                input_fields=self.input_schema.field_names(),
+                input_fields=list(self.input_schema.model_fields),
                 generated_fields=field_names,
                 total_input_tokens=per_record_stats.total_input_tokens,
                 total_output_tokens=per_record_stats.total_output_tokens,
@@ -148,7 +146,7 @@ class ConvertOp(PhysicalOperator, ABC):
         pass
 
     @abstractmethod
-    def convert(self, candidate: DataRecord, fields: dict[str, Field]) -> tuple[dict[str, list], GenerationStats]:
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[str, list], GenerationStats]:
         """
         This abstract method will be implemented by subclasses of ConvertOp to process the input DataRecord
         and generate the value(s) for each of the specified fields. If the convert operator is a one-to-many
@@ -182,7 +180,7 @@ class ConvertOp(PhysicalOperator, ABC):
 
         # execute the convert
         field_answers: dict[str, list]
-        fields = {field: field_type for field, field_type in self.output_schema.field_map().items() if field in fields_to_generate}
+        fields = {field: field_type for field, field_type in self.output_schema.model_fields.items() if field in fields_to_generate}
         field_answers, generation_stats = self.convert(candidate=candidate, fields=fields)
         assert all([field in field_answers for field in fields_to_generate]), "Not all fields were generated!"
 
@@ -235,7 +233,7 @@ class NonLLMConvert(ConvertOp):
             quality=1.0,
         )
 
-    def convert(self, candidate: DataRecord, fields: dict[str, Field]) -> tuple[dict[str, list], GenerationStats]:
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[str, list], GenerationStats]:
         # apply UDF to input record
         start_time = time.time()
         field_answers = {}
@@ -282,18 +280,21 @@ class LLMConvert(ConvertOp):
         self,
         model: Model,
         prompt_strategy: PromptStrategy = PromptStrategy.COT_QA,
+        reasoning_effort: str | None = None,
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.model = model
         self.prompt_strategy = prompt_strategy
+        self.reasoning_effort = reasoning_effort
         if model is not None:
-            self.generator = generator_factory(model, prompt_strategy, self.cardinality, self.verbose)
+            self.generator = Generator(model, prompt_strategy, reasoning_effort, self.api_base, self.cardinality, self.verbose)
 
     def __str__(self):
         op = super().__str__()
         op += f"    Prompt Strategy: {self.prompt_strategy}\n"
+        op += f"    Reasoning Effort: {self.reasoning_effort}\n"
         return op
 
     def get_id_params(self):
@@ -301,6 +302,7 @@ class LLMConvert(ConvertOp):
         id_params = {
             "model": None if self.model is None else self.model.value,
             "prompt_strategy": None if self.prompt_strategy is None else self.prompt_strategy.value,
+            "reasoning_effort": self.reasoning_effort,
             **id_params,
         }
 
@@ -311,6 +313,7 @@ class LLMConvert(ConvertOp):
         op_params = {
             "model": self.model,
             "prompt_strategy": self.prompt_strategy,
+            "reasoning_effort": self.reasoning_effort,
             **op_params,
         }
 
@@ -320,7 +323,7 @@ class LLMConvert(ConvertOp):
         return None if self.model is None else self.model.value
 
     def is_image_conversion(self) -> bool:
-        return self.model in get_vision_models()
+        return self.prompt_strategy.is_image_prompt()
 
     def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
         """
@@ -334,13 +337,16 @@ class LLMConvert(ConvertOp):
         est_num_output_tokens = NAIVE_EST_NUM_OUTPUT_TOKENS
 
         # get est. of conversion time per record from model card;
-        # NOTE: model will only be None for code synthesis, which uses GPT-3.5 as fallback
         model_name = self.model.value if getattr(self, "model", None) is not None else Model.GPT_4o_MINI.value
         model_conversion_time_per_record = MODEL_CARDS[model_name]["seconds_per_output_token"] * est_num_output_tokens
 
         # 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():
+            usd_per_input_token = MODEL_CARDS[model_name]["usd_per_audio_input_token"]
+
         model_conversion_usd_per_record = (
-            MODEL_CARDS[model_name]["usd_per_input_token"] * est_num_input_tokens
+            usd_per_input_token * est_num_input_tokens
             + MODEL_CARDS[model_name]["usd_per_output_token"] * est_num_output_tokens
         )
 
@@ -349,7 +355,7 @@ class LLMConvert(ConvertOp):
         cardinality = selectivity * source_op_cost_estimates.cardinality
 
         # estimate quality of output based on the strength of the model being used
-        quality = (MODEL_CARDS[model_name]["overall"] / 100.0) * source_op_cost_estimates.quality
+        quality = (MODEL_CARDS[model_name]["overall"] / 100.0)
 
         return OperatorCostEstimates(
             cardinality=cardinality,
@@ -361,7 +367,7 @@ class LLMConvert(ConvertOp):
 
 class LLMConvertBonded(LLMConvert):
 
-    def convert(self, candidate: DataRecord, fields: dict[str, Field]) -> tuple[dict[str, list], GenerationStats]:
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[str, list], GenerationStats]:
         # get the set of input fields to use for the convert operation
         input_fields = self.get_input_fields()
 

palimpzest/query/operators/critique_and_refine_convert.py

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 from typing import Any
 
+from pydantic.fields import FieldInfo
+
 from palimpzest.constants import MODEL_CARDS, Model, PromptStrategy
-from palimpzest.core.data.dataclasses import GenerationStats, OperatorCostEstimates
 from palimpzest.core.elements.records import DataRecord
-from palimpzest.query.generators.generators import generator_factory
+from palimpzest.core.models import GenerationStats, OperatorCostEstimates
+from palimpzest.query.generators.generators import Generator
 from palimpzest.query.operators.convert import LLMConvert
 
 # TYPE DEFINITIONS
@@ -35,8 +37,8 @@ class CriticAndRefineConvert(LLMConvert):
             raise ValueError(f"Unsupported prompt strategy: {self.prompt_strategy}")
 
         # create generators
-        self.critic_generator = generator_factory(self.critic_model, self.critic_prompt_strategy, self.cardinality, self.verbose)
-        self.refine_generator = generator_factory(self.refine_model, self.refinement_prompt_strategy, self.cardinality, self.verbose)
+        self.critic_generator = Generator(self.critic_model, self.critic_prompt_strategy, self.reasoning_effort, self.api_base, self.cardinality, self.verbose)
+        self.refine_generator = Generator(self.refine_model, self.refinement_prompt_strategy, self.reasoning_effort, self.api_base, self.cardinality, self.verbose)
 
     def __str__(self):
         op = super().__str__()
@@ -86,7 +88,7 @@ class CriticAndRefineConvert(LLMConvert):
 
         return naive_op_cost_estimates
 
-    def convert(self, candidate: DataRecord, fields: list[str]) -> tuple[dict[FieldName, list[Any]], GenerationStats]:
+    def convert(self, candidate: DataRecord, fields: dict[str, FieldInfo]) -> tuple[dict[FieldName, list[Any]], GenerationStats]:
         # get input fields
         input_fields = self.get_input_fields()