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

palimpzest/query/operators/logical.py

@@ -3,11 +3,13 @@ from __future__ import annotations
 import json
 from typing import Callable
 
+from pydantic import BaseModel
+
 from palimpzest.constants import AggFunc, Cardinality
-from palimpzest.core.data.datareaders import DataReader
+from palimpzest.core.data import context, dataset
 from palimpzest.core.elements.filters import Filter
 from palimpzest.core.elements.groupbysig import GroupBySig
-from palimpzest.core.lib.schemas import Schema
+from palimpzest.core.lib.schemas import Average, Count
 from palimpzest.utils.hash_helpers import hash_for_id
 
 
@@ -16,8 +18,8 @@ class LogicalOperator:
     A logical operator is an operator that operates on Sets.
 
     Right now it can be one of:
-    - BaseScan (scans data from DataReader)
-    - CacheScan (scans cached Set)
+    - BaseScan (scans data from a root Dataset)
+    - ContextScan (loads the context for a root Dataset)
     - FilteredScan (scans input Set and applies filter)
     - ConvertScan (scans input Set and converts it to new Schema)
     - LimitScan (scans up to N records from a Set)
@@ -25,6 +27,8 @@ class LogicalOperator:
     - Aggregate (applies an aggregation on the Set)
     - RetrieveScan (fetches documents from a provided input for a given query)
     - Map (applies a function to each record in the Set without adding any new columns)
+    - ComputeOperator (executes a computation described in natural language)
+    - SearchOperator (executes a search query on the input Context)
 
     Every logical operator must declare the get_logical_id_params() and get_logical_op_params() methods,
     which return dictionaries of parameters that are used to compute the logical op id and to implement
@@ -33,17 +37,21 @@ class LogicalOperator:
 
     def __init__(
         self,
-        output_schema: Schema,
-        input_schema: Schema | None = None,
+        output_schema: type[BaseModel],
+        input_schema: type[BaseModel] | None = None,
+        depends_on: list[str] | None = None,
     ):
+        # TODO: can we eliminate input_schema?
         self.output_schema = output_schema
         self.input_schema = input_schema
+        self.depends_on = [] if depends_on is None else sorted(depends_on)
         self.logical_op_id: str | None = None
+        self.unique_logical_op_id: str | None = None
 
         # compute the fields generated by this logical operator
-        input_field_names = self.input_schema.field_names() if self.input_schema is not None else []
+        input_field_names = list(self.input_schema.model_fields) if self.input_schema is not None else []
         self.generated_fields = sorted(
-            [field_name for field_name in self.output_schema.field_names() if field_name not in input_field_names]
+            [field_name for field_name in self.output_schema.model_fields if field_name not in input_field_names]
         )
 
     def __str__(self) -> str:
@@ -54,12 +62,28 @@ class LogicalOperator:
         return isinstance(other, self.__class__) and all_id_params_match
 
     def copy(self) -> LogicalOperator:
-        return self.__class__(**self.get_logical_op_params())
+        logical_op_copy = self.__class__(**self.get_logical_op_params())
+        logical_op_copy.logical_op_id = self.logical_op_id
+        logical_op_copy.unique_logical_op_id = self.unique_logical_op_id
+        return logical_op_copy
 
     def logical_op_name(self) -> str:
         """Name of the logical operator."""
         return str(self.__class__.__name__)
 
+    def get_unique_logical_op_id(self) -> str:
+        """
+        Get the unique logical operator id for this logical operator.
+        """
+        return self.unique_logical_op_id
+
+    def set_unique_logical_op_id(self, unique_logical_op_id: str) -> None:
+        """
+        Set the unique logical operator id for this logical operator.
+        This is used to uniquely identify the logical operator in the query plan.
+        """
+        self.unique_logical_op_id = unique_logical_op_id
+
     def get_logical_id_params(self) -> dict:
         """
         Returns a dictionary mapping of logical operator parameters which are relevant
@@ -69,6 +93,7 @@ class LogicalOperator:
         NOTE: input_schema and output_schema are not included in the id params because
               they depend on how the Optimizer orders operations.
         """
+        # TODO: should we use `generated_fields` after getting rid of them in PhysicalOperator?
         return {"generated_fields": self.generated_fields}
 
     def get_logical_op_params(self) -> dict:
@@ -78,10 +103,16 @@ class LogicalOperator:
 
         NOTE: Should be overriden by subclasses to include class-specific parameters.
         """
-        return {"input_schema": self.input_schema, "output_schema": self.output_schema}
+        return {
+            "input_schema": self.input_schema,
+            "output_schema": self.output_schema,
+            "depends_on": self.depends_on,
+        }
 
     def get_logical_op_id(self):
         """
+        TODO: turn this into a property?
+
         NOTE: We do not call this in the __init__() method as subclasses may set parameters
               returned by self.get_logical_op_params() after they call to super().__init__().
         """
@@ -119,13 +150,19 @@ class Aggregate(LogicalOperator):
     def __init__(
         self,
         agg_func: AggFunc,
-        target_cache_id: str | None = None,
         *args,
         **kwargs,
     ):
+        if kwargs.get("output_schema") is None:
+            if agg_func == AggFunc.COUNT:
+                kwargs["output_schema"] = Count
+            elif agg_func == AggFunc.AVERAGE:
+                kwargs["output_schema"] = Average
+            else:
+                raise ValueError(f"Unsupported aggregation function: {agg_func}")
+
         super().__init__(*args, **kwargs)
         self.agg_func = agg_func
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
         return f"{self.__class__.__name__}(function: {str(self.agg_func.value)})"
@@ -140,7 +177,6 @@ class Aggregate(LogicalOperator):
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
             "agg_func": self.agg_func,
-            "target_cache_id": self.target_cache_id,
             **logical_op_params,
         }
 
@@ -148,75 +184,87 @@ class Aggregate(LogicalOperator):
 
 
 class BaseScan(LogicalOperator):
-    """A BaseScan is a logical operator that represents a scan of a particular data source."""
+    """A BaseScan is a logical operator that represents a scan of a particular root Dataset."""
 
-    def __init__(self, datareader: DataReader, output_schema: Schema):
-        super().__init__(output_schema=output_schema)
-        self.datareader = datareader
+    def __init__(self, datasource: dataset.Dataset, output_schema: type[BaseModel], *args, **kwargs):
+        super().__init__(*args, output_schema=output_schema, **kwargs)
+        self.datasource = datasource
 
     def __str__(self):
-        return f"BaseScan({self.datareader},{self.output_schema})"
+        return f"BaseScan({self.datasource},{self.output_schema})"
 
     def __eq__(self, other) -> bool:
         return (
             isinstance(other, BaseScan)
-            and self.input_schema.get_desc() == other.input_schema.get_desc()
-            and self.output_schema.get_desc() == other.output_schema.get_desc()
-            and self.datareader == other.datareader
+            and self.input_schema == other.input_schema
+            and self.output_schema == other.output_schema
+            and self.datasource == other.datasource
         )
 
     def get_logical_id_params(self) -> dict:
-        return super().get_logical_id_params()
+        logical_id_params = super().get_logical_id_params()
+        logical_id_params = {
+            "id": self.datasource.id,
+            **logical_id_params,
+        }
+
+        return logical_id_params
 
     def get_logical_op_params(self) -> dict:
         logical_op_params = super().get_logical_op_params()
-        logical_op_params = {"datareader": self.datareader, **logical_op_params}
+        logical_op_params = {"datasource": self.datasource, **logical_op_params}
 
         return logical_op_params
 
 
-class CacheScan(LogicalOperator):
-    """A CacheScan is a logical operator that represents a scan of a cached Set."""
+class ContextScan(LogicalOperator):
+    """A ContextScan is a logical operator that loads the context for a particular root Dataset."""
 
-    def __init__(self, datareader: DataReader, output_schema: Schema):
-        super().__init__(output_schema=output_schema)
-        self.datareader = datareader
+    def __init__(self, context: context.Context, output_schema: type[BaseModel], *args, **kwargs):
+        super().__init__(*args, output_schema=output_schema, **kwargs)
+        self.context = context
 
     def __str__(self):
-        return f"CacheScan({self.datareader},{self.output_schema})"
+        return f"ContextScan({self.context},{self.output_schema})"
+
+    def __eq__(self, other) -> bool:
+        return (
+            isinstance(other, ContextScan)
+            and self.context.id == other.context.id
+        )
 
     def get_logical_id_params(self) -> dict:
-        return super().get_logical_id_params()
+        logical_id_params = super().get_logical_id_params()
+        logical_id_params = {
+            "id": self.context.id,
+            **logical_id_params,
+        }
+
+        return logical_id_params
 
     def get_logical_op_params(self) -> dict:
         logical_op_params = super().get_logical_op_params()
-        logical_op_params = {"datareader": self.datareader, **logical_op_params}
+        logical_op_params = {"context": self.context, **logical_op_params}
 
         return logical_op_params
 
 
 class ConvertScan(LogicalOperator):
-    """A ConvertScan is a logical operator that represents a scan of a particular data source, with conversion applied."""
+    """A ConvertScan is a logical operator that represents a scan of a particular input Dataset, with conversion applied."""
 
     def __init__(
         self,
         cardinality: Cardinality = Cardinality.ONE_TO_ONE,
         udf: Callable | None = None,
-        depends_on: list[str] | None = None,
-        desc: str | None = None,
-        target_cache_id: str | None = None,
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.cardinality = cardinality
         self.udf = udf
-        self.depends_on = [] if depends_on is None else sorted(depends_on)
-        self.desc = desc
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
-        return f"ConvertScan({self.input_schema} -> {str(self.output_schema)},{str(self.desc)})"
+        return f"ConvertScan({self.input_schema} -> {str(self.output_schema)})"
 
     def get_logical_id_params(self) -> dict:
         logical_id_params = super().get_logical_id_params()
@@ -233,9 +281,40 @@ class ConvertScan(LogicalOperator):
         logical_op_params = {
             "cardinality": self.cardinality,
             "udf": self.udf,
-            "depends_on": self.depends_on,
-            "desc": self.desc,
-            "target_cache_id": self.target_cache_id,
+            **logical_op_params,
+        }
+
+        return logical_op_params
+
+
+class Distinct(LogicalOperator):
+    def __init__(self, distinct_cols: list[str] | None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # if distinct_cols is not None, check that all columns are in the input schema
+        if distinct_cols is not None:
+            for col in distinct_cols:
+                assert col in self.input_schema.model_fields, f"Column {col} not found in input schema {self.input_schema} for Distinct operator"
+
+        # store the list of distinct columns, sorted
+        self.distinct_cols = (
+            sorted([field_name for field_name in self.input_schema.model_fields])
+            if distinct_cols is None
+            else sorted(distinct_cols)
+        )
+
+    def __str__(self):
+        return f"Distinct({self.distinct_cols})"
+
+    def get_logical_id_params(self) -> dict:
+        logical_id_params = super().get_logical_id_params()
+        logical_id_params = {"distinct_cols": self.distinct_cols, **logical_id_params}
+
+        return logical_id_params
+
+    def get_logical_op_params(self) -> dict:
+        logical_op_params = super().get_logical_op_params()
+        logical_op_params = {
+            "distinct_cols": self.distinct_cols,
             **logical_op_params,
         }
 
@@ -243,20 +322,16 @@ class ConvertScan(LogicalOperator):
 
 
 class FilteredScan(LogicalOperator):
-    """A FilteredScan is a logical operator that represents a scan of a particular data source, with filters applied."""
+    """A FilteredScan is a logical operator that represents a scan of a particular input Dataset, with filters applied."""
 
     def __init__(
         self,
         filter: Filter,
-        depends_on: list[str] | None = None,
-        target_cache_id: str | None = None,
         *args,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.filter = filter
-        self.depends_on = [] if depends_on is None else sorted(depends_on)
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
         return f"FilteredScan({str(self.output_schema)}, {str(self.filter)})"
@@ -274,8 +349,6 @@ class FilteredScan(LogicalOperator):
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
             "filter": self.filter,
-            "depends_on": self.depends_on,
-            "target_cache_id": self.target_cache_id,
             **logical_op_params,
         }
 
@@ -286,7 +359,6 @@ class GroupByAggregate(LogicalOperator):
     def __init__(
         self,
         group_by_sig: GroupBySig,
-        target_cache_id: str | None = None,
         *args,
         **kwargs,
     ):
@@ -297,7 +369,6 @@ class GroupByAggregate(LogicalOperator):
         if not valid:
             raise TypeError(error)
         self.group_by_sig = group_by_sig
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
         return f"GroupBy({self.group_by_sig.serialize()})"
@@ -312,7 +383,30 @@ class GroupByAggregate(LogicalOperator):
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
             "group_by_sig": self.group_by_sig,
-            "target_cache_id": self.target_cache_id,
+            **logical_op_params,
+        }
+
+        return logical_op_params
+
+
+class JoinOp(LogicalOperator):
+    def __init__(self, condition: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.condition = condition
+
+    def __str__(self):
+        return f"Join(condition={self.condition})"
+
+    def get_logical_id_params(self) -> dict:
+        logical_id_params = super().get_logical_id_params()
+        logical_id_params = {"condition": self.condition, **logical_id_params}
+
+        return logical_id_params
+
+    def get_logical_op_params(self) -> dict:
+        logical_op_params = super().get_logical_op_params()
+        logical_op_params = {
+            "condition": self.condition,
             **logical_op_params,
         }
 
@@ -320,10 +414,9 @@ class GroupByAggregate(LogicalOperator):
 
 
 class LimitScan(LogicalOperator):
-    def __init__(self, limit: int, target_cache_id: str | None = None, *args, **kwargs):
+    def __init__(self, limit: int, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.limit = limit
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
         return f"LimitScan({str(self.input_schema)}, {str(self.output_schema)})"
@@ -338,7 +431,6 @@ class LimitScan(LogicalOperator):
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
             "limit": self.limit,
-            "target_cache_id": self.target_cache_id,
             **logical_op_params,
         }
 
@@ -346,10 +438,9 @@ class LimitScan(LogicalOperator):
 
 
 class Project(LogicalOperator):
-    def __init__(self, project_cols: list[str], target_cache_id: str | None = None, *args, **kwargs):
+    def __init__(self, project_cols: list[str], *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.project_cols = project_cols
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
         return f"Project({self.input_schema}, {self.project_cols})"
@@ -364,7 +455,6 @@ class Project(LogicalOperator):
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
             "project_cols": self.project_cols,
-            "target_cache_id": self.target_cache_id,
             **logical_op_params,
         }
 
@@ -372,7 +462,7 @@ class Project(LogicalOperator):
 
 
 class RetrieveScan(LogicalOperator):
-    """A RetrieveScan is a logical operator that represents a scan of a particular data source, with a convert-like retrieve applied."""
+    """A RetrieveScan is a logical operator that represents a scan of a particular input Dataset, with a convert-like retrieve applied."""
 
     def __init__(
         self,
@@ -381,7 +471,6 @@ class RetrieveScan(LogicalOperator):
         search_attr,
         output_attrs,
         k,
-        target_cache_id: str = None,
         *args,
         **kwargs,
     ):
@@ -391,10 +480,9 @@ class RetrieveScan(LogicalOperator):
         self.search_attr = search_attr
         self.output_attrs = output_attrs
         self.k = k
-        self.target_cache_id = target_cache_id
 
     def __str__(self):
-        return f"RetrieveScan({self.input_schema} -> {str(self.output_schema)},{str(self.desc)})"
+        return f"RetrieveScan({self.input_schema} -> {str(self.output_schema)})"
 
     def get_logical_id_params(self) -> dict:
         # NOTE: if we allow optimization over index, then we will need to include it in the id params
@@ -418,36 +506,31 @@ class RetrieveScan(LogicalOperator):
             "search_attr": self.search_attr,
             "output_attrs": self.output_attrs,
             "k": self.k,
-            "target_cache_id": self.target_cache_id,
             **logical_op_params,
         }
 
         return logical_op_params
 
 
-# TODO: (near-term) maybe we should try to fold this into ConvertScan, and make the internals of PZ
-#       amenable to a convert operator (with a UDF) that does not add new columns?
-class MapScan(LogicalOperator):
-    """A MapScan is a logical operator that applies a UDF to each input record without adding new columns."""
+class ComputeOperator(LogicalOperator):
+    """
+    A ComputeOperator is a logical operator that performs a computation described in natural language
+    on a given Context.
+    """
 
-    def __init__(
-        self,
-        udf: Callable | None = None,
-        target_cache_id: str | None = None,
-        *args,
-        **kwargs,
-    ):
+    def __init__(self, context_id: str, instruction: str, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.udf = udf
-        self.target_cache_id = target_cache_id
+        self.context_id = context_id
+        self.instruction = instruction
 
     def __str__(self):
-        return f"MapScan({self.output_schema}, {self.udf.__name__})"
+        return f"ComputeOperator(id={self.context_id}, instr={self.instruction:20s})"
 
     def get_logical_id_params(self) -> dict:
         logical_id_params = super().get_logical_id_params()
         logical_id_params = {
-            "udf": self.udf,
+            "context_id": self.context_id,
+            "instruction": self.instruction,
             **logical_id_params,
         }
 
@@ -456,8 +539,43 @@ class MapScan(LogicalOperator):
     def get_logical_op_params(self) -> dict:
         logical_op_params = super().get_logical_op_params()
         logical_op_params = {
-            "udf": self.udf,
-            "target_cache_id": self.target_cache_id,
+            "context_id": self.context_id,
+            "instruction": self.instruction,
+            **logical_op_params,
+        }
+
+        return logical_op_params
+
+
+class SearchOperator(LogicalOperator):
+    """
+    A SearchOperator is a logical operator that executes a search described in natural language
+    on a given Context.
+    """
+
+    def __init__(self, context_id: str, search_query: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.context_id = context_id
+        self.search_query = search_query
+
+    def __str__(self):
+        return f"SearchOperator(id={self.context_id}, search_query={self.search_query:20s})"
+
+    def get_logical_id_params(self) -> dict:
+        logical_id_params = super().get_logical_id_params()
+        logical_id_params = {
+            "context_id": self.context_id,
+            "search_query": self.search_query,
+            **logical_id_params,
+        }
+
+        return logical_id_params
+
+    def get_logical_op_params(self) -> dict:
+        logical_op_params = super().get_logical_op_params()
+        logical_op_params = {
+            "context_id": self.context_id,
+            "search_query": self.search_query,
             **logical_op_params,
         }
 

palimpzest/query/operators/mixture_of_agents_convert.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+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.core.lib.fields import Field
-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
@@ -20,7 +21,6 @@ class MixtureOfAgentsConvert(LLMConvert):
         aggregator_model: Model,
         proposer_prompt_strategy: PromptStrategy = PromptStrategy.COT_MOA_PROPOSER,
         aggregator_prompt_strategy: PromptStrategy = PromptStrategy.COT_MOA_AGG,
-        proposer_prompt: str | None = None,
         *args,
         **kwargs,
     ):
@@ -33,14 +33,13 @@ class MixtureOfAgentsConvert(LLMConvert):
         self.aggregator_model = aggregator_model
         self.proposer_prompt_strategy = proposer_prompt_strategy
         self.aggregator_prompt_strategy = aggregator_prompt_strategy
-        self.proposer_prompt = proposer_prompt
 
         # create generators
         self.proposer_generators = [
-            generator_factory(model, self.proposer_prompt_strategy, self.cardinality, self.verbose)
+            Generator(model, self.proposer_prompt_strategy, self.reasoning_effort, self.api_base, self.cardinality, self.verbose)
             for model in proposer_models
         ]
-        self.aggregator_generator = generator_factory(aggregator_model, self.aggregator_prompt_strategy, self.cardinality, self.verbose)
+        self.aggregator_generator = Generator(aggregator_model, self.aggregator_prompt_strategy, self.reasoning_effort, self.api_base, self.cardinality, self.verbose)
 
     def __str__(self):
         op = super().__str__()
@@ -77,6 +76,9 @@ class MixtureOfAgentsConvert(LLMConvert):
 
         return op_params
 
+    def is_image_conversion(self) -> bool:
+        return self.proposer_prompt_strategy.is_image_prompt()
+
     def naive_cost_estimates(self, source_op_cost_estimates: OperatorCostEstimates) -> OperatorCostEstimates:
         """
         Currently, we are using multiple proposer models with different temperatures to synthesize
@@ -111,7 +113,7 @@ class MixtureOfAgentsConvert(LLMConvert):
 
         return naive_op_cost_estimates
 
-    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 input fields
         input_fields = self.get_input_fields()