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

palimpzest/query/processor/query_processor_factory.py

@@ -1,17 +1,17 @@
 import logging
 from enum import Enum
 
+from palimpzest.core.data.dataset import Dataset
 from palimpzest.core.elements.records import DataRecordCollection
 from palimpzest.query.execution.execution_strategy import ExecutionStrategy, SentinelExecutionStrategy
 from palimpzest.query.execution.execution_strategy_type import ExecutionStrategyType, SentinelExecutionStrategyType
-from palimpzest.query.optimizer.cost_model import CostModel
+from palimpzest.query.optimizer.cost_model import SampleBasedCostModel
 from palimpzest.query.optimizer.optimizer import Optimizer
 from palimpzest.query.optimizer.optimizer_strategy_type import OptimizationStrategyType
 from palimpzest.query.processor.config import QueryProcessorConfig
-from palimpzest.query.processor.processing_strategy_type import ProcessingStrategyType
 from palimpzest.query.processor.query_processor import QueryProcessor
-from palimpzest.sets import Dataset, Set
 from palimpzest.utils.model_helpers import get_models
+from palimpzest.validator.validator import Validator
 
 logger = logging.getLogger(__name__)
 
@@ -33,12 +33,11 @@ class QueryProcessorFactory:
         an exception if the conversion fails.
         """
         strategy_types = {
-            "processing_strategy": ProcessingStrategyType,
             "execution_strategy": ExecutionStrategyType,
             "sentinel_execution_strategy": SentinelExecutionStrategyType,
             "optimizer_strategy": OptimizationStrategyType,
         }
-        for strategy in ["processing_strategy", "execution_strategy", "sentinel_execution_strategy", "optimizer_strategy"]:
+        for strategy in ["execution_strategy", "sentinel_execution_strategy", "optimizer_strategy"]:
             strategy_str = getattr(config, strategy)
             strategy_type = strategy_types[strategy]
             strategy_enum = None
@@ -54,57 +53,62 @@ class QueryProcessorFactory:
         return config
 
     @classmethod
-    def _config_validation_and_normalization(cls, config: QueryProcessorConfig):
+    def _config_validation_and_normalization(cls, config: QueryProcessorConfig, train_dataset: dict[str, Dataset] | None, validator : Validator | None):
         if config.policy is None:
             raise ValueError("Policy is required for optimizer")
-
-        if config.cache:
-            raise ValueError("cache=True is not supported yet")
         
         # only one of progress or verbose can be set; we will default to progress=True
         if config.progress and config.verbose:
             print("WARNING: Both `progress` and `verbose` are set to True, but only one can be True at a time; defaulting to `progress=True`")
             config.verbose = False
 
-        # handle "auto" defaults for processing and sentinel execution strategies
-        if config.processing_strategy == "auto":
-            config.processing_strategy = "no_sentinel" if config.val_datasource is None else "sentinel"
+        # boolean flag for whether we're performing optimization or not
+        optimization = train_dataset is not None or validator is not None
+        val_based_opt = train_dataset is None and validator is not None
 
+        # handle "auto" default for sentinel execution strategies
         if config.sentinel_execution_strategy == "auto":
-            config.sentinel_execution_strategy = None if config.val_datasource is None else "mab"
+            config.sentinel_execution_strategy = ("validator" if val_based_opt else "mab") if optimization else None
 
         # convert the config values for processing, execution, and optimization strategies to enums
         config = cls._normalize_strategies(config)
 
-        # check that processor uses a supported execution strategy
-        if config.execution_strategy not in config.processing_strategy.valid_execution_strategies():
-            raise ValueError(f"Unsupported `execution_strategy` {config.execution_strategy} for `processing_strategy` {config.processing_strategy}.")
-
-        # check that validation data is provided for sentinel execution
-        if config.val_datasource is None and config.processing_strategy.is_sentinel_strategy():
-            raise ValueError("`val_datasource` is required for SENTINEL processing strategies")
-
-        # check that sentinel execution is provided for sentinel processor
-        if config.sentinel_execution_strategy is None and config.processing_strategy.is_sentinel_strategy():
-            raise ValueError("`sentinel_execution_strategy` is required for SENTINEL processing strategies")
-
         # get available models
         available_models = getattr(config, 'available_models', [])
         if available_models is None or len(available_models) == 0:
-            available_models = get_models(include_vision=True)
+            available_models = get_models(use_vertex=config.use_vertex, gemini_credentials_path=config.gemini_credentials_path, api_base=config.api_base)
+
+        # remove any models specified in the config
+        remove_models = getattr(config, 'remove_models', [])
+        if remove_models is not None and len(remove_models) > 0:
+            available_models = [model for model in available_models if model not in remove_models]
+            logger.info(f"Removed models from available models based on config: {remove_models}")
+
+        # set the final set of available models in the config
         config.available_models = available_models
 
         return config
 
     @classmethod
     def _create_optimizer(cls, config: QueryProcessorConfig) -> Optimizer:
-        return Optimizer(cost_model=CostModel(), **config.to_dict())
+        return Optimizer(cost_model=SampleBasedCostModel(), **config.to_dict())
 
     @classmethod
-    def _create_execution_strategy(cls, config: QueryProcessorConfig) -> ExecutionStrategy:
+    def _create_execution_strategy(cls, dataset: Dataset, config: QueryProcessorConfig) -> ExecutionStrategy:
         """
         Creates an execution strategy based on the configuration.
         """
+        # for parallel execution, set the batch size if there's a limit in the query
+        limit = dataset.get_limit()
+        if limit is not None and config.execution_strategy == ExecutionStrategyType.PARALLEL:
+            if config.batch_size is None:
+                config.batch_size = limit
+                logger.info(f"Setting batch size to query limit: {limit}")
+            elif config.batch_size > limit:
+                config.batch_size = limit
+                logger.info(f"Setting batch size to query limit: {limit} since it was larger than the limit")
+
+        # create the execution strategy
         execution_strategy_cls = config.execution_strategy.value
         return execution_strategy_cls(**config.to_dict())
 
@@ -122,9 +126,10 @@ class QueryProcessorFactory:
     @classmethod
     def create_processor(
         cls,
-        dataset: Set,
+        dataset: Dataset,
         config: QueryProcessorConfig | None = None,
-        **kwargs
+        train_dataset: dict[str, Dataset] | None = None,
+        validator: Validator | None = None,
     ) -> QueryProcessor:
         """
         Creates a QueryProcessor with specified processing and execution strategies.
@@ -138,23 +143,26 @@ class QueryProcessorFactory:
             config = QueryProcessorConfig()
 
         # apply any additional keyword arguments to the config and validate its contents
-        config.update(**kwargs)
-        config = cls._config_validation_and_normalization(config)
+        config = cls._config_validation_and_normalization(config, train_dataset, validator)
 
         # create the optimizer, execution strateg(ies), and processor
         optimizer = cls._create_optimizer(config)
-        config.execution_strategy = cls._create_execution_strategy(config)
+        config.execution_strategy = cls._create_execution_strategy(dataset, config)
         config.sentinel_execution_strategy = cls._create_sentinel_execution_strategy(config)
-        processor_cls = config.processing_strategy.value
-        processor = processor_cls(dataset, optimizer, **config.to_dict())
+        processor = QueryProcessor(dataset, optimizer, train_dataset=train_dataset, validator=validator, **config.to_dict())
 
         return processor
 
     @classmethod
-    def create_and_run_processor(cls, dataset: Dataset, config: QueryProcessorConfig | None = None, **kwargs) -> DataRecordCollection:
-        # TODO(Jun): Consider to use cache here.
+    def create_and_run_processor(
+        cls,
+        dataset: Dataset,
+        config: QueryProcessorConfig | None = None,
+        train_dataset: dict[str, Dataset] | None = None,
+        validator: Validator | None = None,
+    ) -> DataRecordCollection:
         logger.info(f"Creating processor for dataset: {dataset}")
-        processor = cls.create_processor(dataset=dataset, config=config, **kwargs)
+        processor = cls.create_processor(dataset, config, train_dataset, validator)
         logger.info(f"Created processor: {processor}")
 
         return processor.execute()

palimpzest/schemabuilder/schema_builder.py

@@ -7,15 +7,15 @@ This method is a simple wrapper for different methods, e.g., from_csv, from_yml,
 
 import json
 import os
+from typing import Any
 
 import pandas as pd
 import yaml
+from pydantic import BaseModel
 from pyld import jsonld
 
-import palimpzest.core.lib.fields as pz_fields
 import palimpzest.core.lib.schemas as pz_schemas
-from palimpzest.core.lib.fields import Field
-from palimpzest.core.lib.schemas import Schema
+from palimpzest.core.lib.schemas import create_schema_from_fields
 
 
 class SchemaBuilder:
@@ -24,19 +24,17 @@ class SchemaBuilder:
     def from_file(cls,
         schema_file: str,
         schema_name: str = "",
-        schema_description: str = "",
         include_attributes: list = None,
         exclude_attributes: list = None,
-        schema_type: Schema = None,
+        schema_type: BaseModel = None,
         ):
         """
         Inputs:
             schema_file: str - the path to the file
-            description (optional): str - the description of the schema
             name (optional): str - the name of the schema
             include_attributes (optional): list - a list of attribute names to include in the schema. If None, all attributes are included.
             exclude_attributes (optional): list - a list of attribute names to exclude from the schema. If None, no attributes are excluded.
-            schema_type (optional): Schema - the parent type of the schema to generate, e.g. ScientificPapers have a schema_type of PDFFile. If None, a generic Schema type is used.
+            schema_type (optional): BaseModel - the parent type of the schema to generate, e.g. ScientificPapers have a schema_type of PDFFile. If None, a generic Schema type is used.
         Outputs:
             A class object - the dynamically generated class
         """
@@ -64,12 +62,6 @@ class SchemaBuilder:
             else:
                 schema_name = "".join([word.capitalize() for word in basename.split("_")])
 
-        if not schema_description:
-            if schema_data['description']:
-                schema_description = schema_data['description']
-            else:
-                schema_description = f"A schema generated from the {file_extension} file {basename}."
-
         if include_attributes is None:
            include_attributes = []
         if exclude_attributes is None:
@@ -78,14 +70,16 @@ class SchemaBuilder:
         if schema_type is None:
             if schema_data.get('type', None):
                 # Find if the schema type is a valid class in pz
-                parsed_type = getattr(pz_schemas
-                                      , schema_data['type'], Schema)
-                schema_type = parsed_type if issubclass(parsed_type, Schema) else Schema
+                parsed_type = getattr(pz_schemas, schema_data['type'], BaseModel)
+                schema_type = parsed_type if issubclass(parsed_type, BaseModel) else BaseModel
             else:
-                schema_type = Schema
+                schema_type = BaseModel
 
         # Generate the schema class dynamically
-        attributes = {"__doc__": schema_description}
+        fields = [
+            {"name": field_name, "description": field.description, "type": field.annotation}
+            for field_name, field in schema_type.model_fields.items()
+        ]
         include_attributes_lower = set([a.lower() for a in include_attributes])
         exclude_attributes_lower = set([a.lower() for a in exclude_attributes])
         for field in schema_data['fields']:
@@ -96,15 +90,9 @@ class SchemaBuilder:
                 continue
             name = field['name']
             description = field.get('description', '')
-            field_type = field.get('type', 'Field')
-            field_type = getattr(pz_fields, field_type, Field)
-            if not issubclass(field_type, Field):
-                field_type = Field
-                  
-            attributes[name] = field_type(desc=description)
+            fields.append({"name": name, "description": description, "type": Any})
 
-        # Create the class dynamically
-        return type(schema_name, (schema_type,), attributes)
+        return create_schema_from_fields(fields)
 
     @classmethod
     def from_csv(
@@ -189,9 +177,8 @@ class SchemaBuilder:
             cls,
             schema_file: str,
             schema_name: str = "",
-            schema_description: str = "",
             include_attributes: list = None,
-            schema_type: Schema = None,
+            schema_type: BaseModel = None,
     )-> dict:
         """
         The attributes are extracted from the JSON objects.

palimpzest/utils/model_helpers.py

@@ -3,29 +3,8 @@ import os
 from palimpzest.constants import Model
 
 
-def get_vision_models() -> list[Model]:
-    """
-    Return the set of vision models which the system has access to based on the set of environment variables.
-    """
-    models = []
-    if os.getenv("OPENAI_API_KEY") is not None:
-        openai_vision_models = [
-            model for model in Model
-            if model.is_openai_model() and model.is_vision_model()
-        ]
-        models.extend(openai_vision_models)
-
-    if os.getenv("TOGETHER_API_KEY") is not None:
-        together_vision_models = [
-            model for model in Model
-            if model.is_together_model() and model.is_vision_model()
-        ]
-        models.extend(together_vision_models)
-
-    return models
-
-
-def get_models(include_vision: bool = False, include_embedding: bool = False) -> list[Model]:
+# TODO: better handle vertex vs. google for gemini models
+def get_models(include_embedding: bool = False, use_vertex: bool = True, gemini_credentials_path: str | None = None, api_base: str | None = None) -> list[Model]:
     """
     Return the set of models which the system has access to based on the set environment variables.
     """
@@ -40,67 +19,43 @@ def get_models(include_vision: bool = False, include_embedding: bool = False) ->
 
     if os.getenv("TOGETHER_API_KEY") is not None:
         together_models = [model for model in Model if model.is_together_model()]
-        if not include_vision:
-            together_models = [
-                model for model in together_models if not model.is_vision_model()
-            ]
         if not include_embedding:
             together_models = [
                 model for model in together_models if not model.is_embedding_model()
             ]
         models.extend(together_models)
 
-    if include_vision:
-        vision_models = get_vision_models()
-        models.extend(vision_models)
-
-    return models
-
-# The order is the priority of the model
-TEXT_MODEL_PRIORITY = [
-    # Model.o1,
-    Model.GPT_4o,
-    Model.GPT_4o_MINI,
-    Model.LLAMA3_3_70B,
-    Model.MIXTRAL,
-    Model.DEEPSEEK_V3,
-    Model.LLAMA3_2_3B,
-    Model.LLAMA3_1_8B,
-    Model.DEEPSEEK_R1_DISTILL_QWEN_1_5B,
-]
-
-VISION_MODEL_PRIORITY = [
-    Model.GPT_4o,
-    Model.GPT_4o_MINI,
-    Model.LLAMA3_2_90B_V,
-]
-def get_champion_model(available_models, vision=False):
-    # Select appropriate priority list based on task
-    model_priority = VISION_MODEL_PRIORITY if vision else TEXT_MODEL_PRIORITY
-
-    # Return first available model from priority list
-    for model in model_priority:
-        if model in available_models:
-            return model
+    if os.getenv("ANTHROPIC_API_KEY") is not None:
+        anthropic_models = [model for model in Model if model.is_anthropic_model()]
+        if not include_embedding:
+            anthropic_models = [
+                model for model in anthropic_models if not model.is_embedding_model()
+            ]
+        models.extend(anthropic_models)
 
-    # If no suitable model found, raise informative error
-    task_type = "vision" if vision else "text"
-    raise Exception(
-        f"No {task_type} models available to create physical plans!\n"
-        "You must set at least one of the following environment variables:\n"
-        "[OPENAI_API_KEY, TOGETHER_API_KEY, GOOGLE_API_KEY]\n"
-        f"Available models: {available_models}"
+    gemini_credentials_path = (
+        os.path.join(os.path.expanduser("~"), ".config", "gcloud", "application_default_credentials.json")
+        if gemini_credentials_path is None
+        else gemini_credentials_path
     )
+    if os.getenv("GEMINI_API_KEY") is not None or os.path.exists(gemini_credentials_path):
+        vertex_models = [model for model in Model if model.is_vertex_model()]
+        google_models = [model for model in Model if model.is_google_model()]
+        if not include_embedding:
+            vertex_models = [
+                model for model in vertex_models if not model.is_embedding_model()
+            ]
+        if use_vertex:
+            models.extend(vertex_models)
+        else:
+            models.extend(google_models)
 
+    if api_base is not None:
+        vllm_models = [model for model in Model if model.is_vllm_model()]
+        if not include_embedding:
+            vllm_models = [
+                model for model in vllm_models if not model.is_embedding_model()
+            ]
+        models.extend(vllm_models)
 
-def get_fallback_model(available_models, vision=False):
-    return get_champion_model(available_models, vision)
-
-
-def get_code_champion_model(available_models):
-    # NOTE: for now, assume same champion as get_champion_model()
-    return get_champion_model(available_models, vision=False)
-
-
-def get_champion_model_name(available_models, vision=False):
-    return get_champion_model(available_models, vision).value
+    return models