palimpzest 1.1.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

palimpzest/constants.py

@@ -31,9 +31,9 @@ class Model(str, Enum):
     GEMINI_2_0_FLASH = "vertex_ai/gemini-2.0-flash"
     GEMINI_2_5_FLASH = "vertex_ai/gemini-2.5-flash"
     GEMINI_2_5_PRO = "vertex_ai/gemini-2.5-pro"
-    GOOGLE_GEMINI_2_5_FLASH = "google/gemini-2.5-flash"
-    GOOGLE_GEMINI_2_5_FLASH_LITE = "google/gemini-2.5-flash-lite"
-    GOOGLE_GEMINI_2_5_PRO = "google/gemini-2.5-pro"
+    GOOGLE_GEMINI_2_5_FLASH = "gemini/gemini-2.5-flash"
+    GOOGLE_GEMINI_2_5_FLASH_LITE = "gemini/gemini-2.5-flash-lite"
+    GOOGLE_GEMINI_2_5_PRO = "gemini/gemini-2.5-pro"
     LLAMA_4_MAVERICK = "vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas"
     GPT_4o_AUDIO_PREVIEW = "openai/gpt-4o-audio-preview"
     GPT_4o_MINI_AUDIO_PREVIEW = "openai/gpt-4o-mini-audio-preview"
@@ -72,8 +72,8 @@ class Model(str, Enum):
     def is_vertex_model(self):
         return "vertex_ai" in self.value.lower()
 
-    def is_google_model(self):
-        return "google" in self.value.lower()
+    def is_google_ai_studio_model(self):
+        return "gemini/" in self.value.lower()
 
     def is_vllm_model(self):
         return "hosted_vllm" in self.value.lower()

palimpzest/core/models.py

@@ -35,12 +35,18 @@ class GenerationStats(BaseModel):
     # typed as a float because GenerationStats may be amortized (i.e. divided) across a number of output records
     total_output_tokens: float = 0.0
 
+    # the total number of input tokens processed by embedding models
+    total_embedding_input_tokens: float = 0.0
+
     # the total cost of processing the input tokens; None if this operation did not use an LLM
     total_input_cost: float = 0.0
 
     # the total cost of processing the output tokens; None if this operation did not use an LLM
     total_output_cost: float = 0.0
 
+    # the total cost of processing input tokens for embedding models
+    total_embedding_cost: float = 0.0
+
     # the total cost of processing the input and output tokens; None if this operation did not use an LLM
     cost_per_record: float = 0.0
 
@@ -68,6 +74,9 @@ class GenerationStats(BaseModel):
             "fn_call_duration_secs",
             "total_llm_calls",
             "total_embedding_llm_calls",
+            "total_embedding_input_tokens",
+            "total_embedding_cost"
+            
         ]:
             setattr(self, model_field, getattr(self, model_field) + getattr(other, model_field))
         return self
@@ -85,6 +94,8 @@ class GenerationStats(BaseModel):
                 "cost_per_record",
                 "total_llm_calls",
                 "total_embedding_llm_calls",
+                "total_embedding_input_tokens",
+                "total_embedding_cost"
             ]
         }
         # dct['raw_answers'] = self.raw_answers + other.raw_answers
@@ -107,6 +118,8 @@ class GenerationStats(BaseModel):
             "fn_call_duration_secs",
             "total_llm_calls",
             "total_embedding_llm_calls",
+            "total_embedding_input_tokens",
+            "total_embedding_cost"
         ]:
             setattr(self, model_field, getattr(self, model_field) / quotient)
         return self
@@ -128,6 +141,8 @@ class GenerationStats(BaseModel):
                 "total_llm_calls",
                 "total_embedding_llm_calls",
                 "cost_per_record",
+                "total_embedding_input_tokens",
+                "total_embedding_cost"
             ]
         }
         dct["model_name"] = self.model_name
@@ -217,6 +232,10 @@ class RecordOpStats(BaseModel):
     # typed as a float because GenerationStats may be amortized (i.e. divided) across a number of output records
     total_output_tokens: float = 0.0
 
+    # the total number of input tokens processed by embedding models
+    # typed as a float because GenerationStats may be amortized (i.e. divided) across a number of output records
+    total_embedding_input_tokens: float = 0.0
+
     # the total cost of processing the input tokens; None if this operation did not use an LLM
     total_input_cost: float = 0.0
 
@@ -278,6 +297,9 @@ class OperatorStats(BaseModel):
     # the total output tokens processed by this operation
     total_output_tokens: int = 0
 
+    #the total embedding input tokens processed by this operation
+    total_embedding_input_tokens: int = 0
+
     # a list of RecordOpStats processed by the operation
     record_op_stats_lst: list[RecordOpStats] = Field(default_factory=list)
 
@@ -309,6 +331,7 @@ class OperatorStats(BaseModel):
             self.total_op_cost += stats.total_op_cost
             self.total_input_tokens += stats.total_input_tokens
             self.total_output_tokens += stats.total_output_tokens
+            self.total_embedding_input_tokens += stats.total_embedding_input_tokens
             self.record_op_stats_lst.extend(stats.record_op_stats_lst)
 
         elif isinstance(stats, RecordOpStats):
@@ -319,6 +342,7 @@ class OperatorStats(BaseModel):
             self.total_op_cost += stats.cost_per_record
             self.total_input_tokens += stats.total_input_tokens
             self.total_output_tokens += stats.total_output_tokens
+            self.total_embedding_input_tokens += stats.total_embedding_input_tokens
 
         else:
             raise TypeError(f"Cannot add {type(stats)} to OperatorStats")
@@ -370,6 +394,9 @@ class BasePlanStats(BaseModel):
     # total output tokens processed by this plan
     total_output_tokens: int = 0
 
+    # total embedding input tokens processed by this plan
+    total_embedding_input_tokens: int = 0
+
     # start time for the plan execution; should be set by calling PlanStats.start()
     start_time: float | None = None
 
@@ -385,6 +412,7 @@ class BasePlanStats(BaseModel):
         self.total_plan_cost = self.sum_op_costs() + self.sum_validation_costs()
         self.total_input_tokens = self.sum_input_tokens() + self.sum_validation_input_tokens()
         self.total_output_tokens = self.sum_output_tokens() + self.sum_validation_output_tokens()
+        self.total_embedding_input_tokens = self.sum_embedding_input_tokens() + self.sum_validation_embedding_input_tokens()
 
     @staticmethod
     @abstractmethod
@@ -415,6 +443,13 @@ class BasePlanStats(BaseModel):
         """
         pass
 
+    @abstractmethod
+    def sum_embedding_input_tokens(self) -> int:
+        """
+        Sum the input embedding tokens processed by all operators in this plan.
+        """
+        pass
+
     @abstractmethod
     def add_record_op_stats(self, unique_full_op_id: str, record_op_stats: RecordOpStats | list[RecordOpStats]) -> None:
         """
@@ -453,6 +488,12 @@ class BasePlanStats(BaseModel):
         Sum the output tokens processed by all validation generations in this plan.
         """
         return sum([gen_stats.total_output_tokens for _, gen_stats in self.validation_gen_stats.items()])
+    
+    def sum_validation_embedding_input_tokens(self) -> int:
+        """
+        Sum the input embedding tokens processed by all validation generations in this plan.
+        """
+        return sum([gen_stats.total_embedding_input_tokens for _, gen_stats in self.validation_gen_stats.items()])
 
     def get_total_cost_so_far(self) -> float:
         """
@@ -501,6 +542,12 @@ class PlanStats(BasePlanStats):
         Sum the output tokens processed by all operators in this plan.
         """
         return sum([op_stats.total_output_tokens for _, op_stats in self.operator_stats.items()])
+    
+    def sum_embedding_input_tokens(self) -> int:
+        """
+        Sum the input embedding tokens processed by all operators in this plan.
+        """
+        return sum([op_stats.total_embedding_input_tokens for _, op_stats in self.operator_stats.items()])
 
     def add_record_op_stats(self, unique_full_op_id: str, record_op_stats: RecordOpStats | list[RecordOpStats]) -> None:
         """
@@ -528,6 +575,7 @@ class PlanStats(BasePlanStats):
         self.total_plan_cost += plan_stats.total_plan_cost
         self.total_input_tokens += plan_stats.total_input_tokens
         self.total_output_tokens += plan_stats.total_output_tokens
+        self.total_embedding_input_tokens += plan_stats.total_embedding_input_tokens
         for unique_full_op_id, op_stats in plan_stats.operator_stats.items():
             if unique_full_op_id in self.operator_stats:
                 self.operator_stats[unique_full_op_id] += op_stats
@@ -539,6 +587,7 @@ class PlanStats(BasePlanStats):
         stats += f"total_plan_cost={self.total_plan_cost} \n"
         stats += f"total_input_tokens={self.total_input_tokens} \n"
         stats += f"total_output_tokens={self.total_output_tokens} \n"
+        stats += f"total_embedding_input_tokens={self.total_embedding_input_tokens} \n"
         for idx, op_stats in enumerate(self.operator_stats.values()):
             stats += f"{idx}. {op_stats.op_name} time={op_stats.total_op_time} cost={op_stats.total_op_cost} \n"
         return stats
@@ -586,6 +635,12 @@ class SentinelPlanStats(BasePlanStats):
         Sum the output tokens processed by all operators in this plan.
         """
         return sum(sum([op_stats.total_output_tokens for _, op_stats in phys_op_stats.items()]) for _, phys_op_stats in self.operator_stats.items())
+    
+    def sum_embedding_input_tokens(self) -> int:
+        """
+        Sum the output tokens processed by all operators in this plan.
+        """
+        return sum(sum([op_stats.total_embedding_input_tokens for _, op_stats in phys_op_stats.items()]) for _, phys_op_stats in self.operator_stats.items())
 
     def add_record_op_stats(self, unique_logical_op_id: str, record_op_stats: RecordOpStats | list[RecordOpStats]) -> None:
         """
@@ -627,6 +682,7 @@ class SentinelPlanStats(BasePlanStats):
         self.total_plan_cost += plan_stats.total_plan_cost
         self.total_input_tokens += plan_stats.total_input_tokens
         self.total_output_tokens += plan_stats.total_output_tokens
+        self.total_embedding_input_tokens += plan_stats.total_embedding_input_tokens
         for unique_logical_op_id, physical_op_stats in plan_stats.operator_stats.items():
             for full_op_id, op_stats in physical_op_stats.items():
                 if unique_logical_op_id in self.operator_stats:
@@ -648,6 +704,7 @@ class SentinelPlanStats(BasePlanStats):
         stats += f"total_plan_cost={self.total_plan_cost} \n"
         stats += f"total_input_tokens={self.total_input_tokens} \n"
         stats += f"total_output_tokens={self.total_output_tokens} \n"
+        stats += f"total_embedding_input_tokens={self.total_embedding_input_tokens} \n"
         for outer_idx, physical_op_stats in enumerate(self.operator_stats.values()):
             total_time = sum([op_stats.total_op_time for op_stats in physical_op_stats.values()])
             total_cost = sum([op_stats.total_op_cost for op_stats in physical_op_stats.values()])
@@ -695,6 +752,9 @@ class ExecutionStats(BaseModel):
     # total number of output tokens processed
     total_output_tokens: int = 0
 
+     # total number of embedding input tokens processed
+    total_embedding_input_tokens: int = 0
+
     # total number of tokens processed
     total_tokens: int = 0
 
@@ -748,7 +808,8 @@ class ExecutionStats(BaseModel):
         # compute the tokens for total execution
         self.total_input_tokens = self.sum_input_tokens()
         self.total_output_tokens = self.sum_output_tokens()
-        self.total_tokens = self.total_input_tokens + self.total_output_tokens
+        self.total_embedding_input_tokens = self.sum_embedding_input_tokens()
+        self.total_tokens = self.total_input_tokens + self.total_output_tokens + self.total_embedding_input_tokens
 
         # compute plan_strs
         self.plan_strs = {plan_id: plan_stats.plan_str for plan_id, plan_stats in self.plan_stats.items()}
@@ -780,6 +841,15 @@ class ExecutionStats(BaseModel):
         sentinel_plan_output_tokens = sum([plan_stats.sum_output_tokens() for _, plan_stats in self.sentinel_plan_stats.items()])
         plan_output_tokens = sum([plan_stats.sum_output_tokens() for _, plan_stats in self.plan_stats.items()])
         return plan_output_tokens + sentinel_plan_output_tokens
+    
+
+    def sum_embedding_input_tokens(self) -> int:
+        """
+        Sum the embedding input tokens processed in this execution
+        """
+        sentinel_plan_embedding_input_tokens = sum([plan_stats.sum_embedding_input_tokens() for _, plan_stats in self.sentinel_plan_stats.items()])
+        plan_embedding_input_tokens = sum([plan_stats.sum_embedding_input_tokens() for _, plan_stats in self.plan_stats.items()])
+        return plan_embedding_input_tokens + sentinel_plan_embedding_input_tokens
 
     def add_plan_stats(self, plan_stats: PlanStats | SentinelPlanStats | list[PlanStats] | list[SentinelPlanStats]) -> None:
         """

palimpzest/query/execution/mab_execution_strategy.py

@@ -777,7 +777,7 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
 
                 # if the operator is a non-llm filter which has filtered out records, remove those records from
                 # all downstream operators' full_op_id_to_sources_not_processed
-                if isinstance(op_set[0], NonLLMFilter):
+                if isinstance(op_set[0], NonLLMFilter) and next_unique_logical_op_id is not None:
                     self._remove_filtered_records_from_downstream_ops(topo_idx, plan, op_frontiers, source_indices_to_all_record_sets)
 
         # finalize plan stats

palimpzest/query/operators/convert.py

@@ -121,8 +121,10 @@ class ConvertOp(PhysicalOperator, ABC):
                 generated_fields=field_names,
                 total_input_tokens=per_record_stats.total_input_tokens,
                 total_output_tokens=per_record_stats.total_output_tokens,
+                total_embedding_input_tokens=per_record_stats.total_embedding_input_tokens,
                 total_input_cost=per_record_stats.total_input_cost,
                 total_output_cost=per_record_stats.total_output_cost,
+                total_embedding_cost=per_record_stats.total_embedding_cost,
                 llm_call_duration_secs=per_record_stats.llm_call_duration_secs,
                 fn_call_duration_secs=per_record_stats.fn_call_duration_secs,
                 total_llm_calls=per_record_stats.total_llm_calls,

palimpzest/query/operators/filter.py

@@ -89,8 +89,10 @@ class FilterOp(PhysicalOperator, ABC):
             filter_str=self.filter_obj.get_filter_str(),
             total_input_tokens=generation_stats.total_input_tokens,
             total_output_tokens=generation_stats.total_output_tokens,
+            total_embedding_input_tokens=generation_stats.total_embedding_input_tokens,
             total_input_cost=generation_stats.total_input_cost,
             total_output_cost=generation_stats.total_output_cost,
+            total_embedding_cost=generation_stats.total_embedding_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,

palimpzest/query/operators/join.py

@@ -376,8 +376,10 @@ class LLMJoin(JoinOp):
             join_condition=self.condition,
             total_input_tokens=generation_stats.total_input_tokens,
             total_output_tokens=generation_stats.total_output_tokens,
+            total_embedding_input_tokens=generation_stats.total_embedding_input_tokens,
             total_input_cost=generation_stats.total_input_cost,
             total_output_cost=generation_stats.total_output_cost,
+            total_embedding_cost=generation_stats.total_embedding_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,
@@ -584,13 +586,13 @@ class EmbeddingJoin(LLMJoin):
             return np.zeros((0, 512)), GenerationStats()
 
         start_time = time.time()
-        total_input_tokens = 0
+        total_embedding_input_tokens = 0
         embeddings = None
         if self.text_only:
             client = OpenAI()
             inputs = [dr.to_json_str(bytes_to_str=True, project_cols=input_fields, sorted=True) for dr in candidates]
             response = client.embeddings.create(input=inputs, model=self.embedding_model.value)
-            total_input_tokens = response.usage.total_tokens
+            total_embedding_input_tokens = response.usage.total_tokens
             embeddings = np.array([item.embedding for item in response.data])
         else:
             model = self.locks.get_model(self.embedding_model.value)
@@ -616,14 +618,16 @@ class EmbeddingJoin(LLMJoin):
 
         # compute cost of embedding(s)
         model_card = MODEL_CARDS[self.embedding_model.value]
-        total_input_cost = model_card["usd_per_input_token"] * total_input_tokens
+        total_embedding_cost = model_card["usd_per_input_token"] * total_embedding_input_tokens
         embedding_gen_stats = GenerationStats(
             model_name=self.embedding_model.value,
-            total_input_tokens=total_input_tokens,
+            total_input_tokens=0.0,
             total_output_tokens=0.0,
-            total_input_cost=total_input_cost,
+            total_embedding_input_tokens=total_embedding_input_tokens,
+            total_input_cost=0.0,
             total_output_cost=0.0,
-            cost_per_record=total_input_cost,
+            total_embedding_cost=total_embedding_cost,
+            cost_per_record=total_embedding_cost,
             llm_call_duration_secs=time.time() - start_time,
             total_llm_calls=1,
             total_embedding_llm_calls=len(candidates),

palimpzest/query/operators/rag.py

@@ -109,15 +109,17 @@ class RAGConvert(LLMConvert):
 
         # compute the generation stats object
         model_card = MODEL_CARDS[model_name]
-        total_input_tokens = response.usage.total_tokens
-        total_input_cost = model_card["usd_per_input_token"] * total_input_tokens
+        total_embedding_input_tokens = response.usage.total_tokens
+        total_embedding_cost = model_card["usd_per_input_token"] * total_embedding_input_tokens
         embed_stats = GenerationStats(
             model_name=model_name,  # NOTE: this should be overwritten by generation model in convert()
-            total_input_tokens=total_input_tokens,
+            total_input_tokens=0.0,
             total_output_tokens=0.0,
-            total_input_cost=total_input_cost,
+            total_embedding_input_tokens=total_embedding_input_tokens,
+            total_input_cost=0.0,
             total_output_cost=0.0,
-            cost_per_record=total_input_cost,
+            total_embedding_cost=total_embedding_cost,
+            cost_per_record=total_embedding_cost,
             llm_call_duration_secs=total_time,
             total_llm_calls=1,
             total_embedding_llm_calls=1,
@@ -318,15 +320,17 @@ class RAGFilter(LLMFilter):
 
         # compute the generation stats object
         model_card = MODEL_CARDS[model_name]
-        total_input_tokens = response.usage.total_tokens
-        total_input_cost = model_card["usd_per_input_token"] * total_input_tokens
+        total_embedding_input_tokens = response.usage.total_tokens
+        total_embedding_cost = model_card["usd_per_input_token"] * total_embedding_input_tokens
         embed_stats = GenerationStats(
             model_name=model_name,  # NOTE: this should be overwritten by generation model in filter()
-            total_input_tokens=total_input_tokens,
+            total_input_tokens=0.0,
             total_output_tokens=0.0,
-            total_input_cost=total_input_cost,
+            total_embedding_input_tokens=total_embedding_input_tokens,
+            total_input_cost=0.0,
             total_output_cost=0.0,
-            cost_per_record=total_input_cost,
+            total_embedding_cost=total_embedding_cost,
+            cost_per_record=total_embedding_cost,
             llm_call_duration_secs=total_time,
             total_llm_calls=1,
             total_embedding_llm_calls=1,

palimpzest/query/optimizer/cost_model.py

@@ -105,9 +105,10 @@ class SampleBasedCostModel:
                     "time_per_record": record_op_stats.time_per_record,
                     "quality": record_op_stats.quality,
                     "passed_operator": record_op_stats.passed_operator,
-                    "source_indices": record_op_stats.record_source_indices,  # TODO: remove
-                    "op_details": record_op_stats.op_details,                 # TODO: remove
-                    "answer": record_op_stats.answer,                         # TODO: remove
+                    "source_indices": record_op_stats.record_source_indices,
+                    "op_details": record_op_stats.op_details,
+                    "answer": record_op_stats.answer,
+                    "op_name": record_op_stats.op_name,
                 }
                 execution_record_op_stats.append(record_op_stats_dict)
 
@@ -128,8 +129,12 @@ class SampleBasedCostModel:
                     else physical_op_df.source_indices.apply(tuple).nunique()
                 )
 
-                # compute selectivity
+                # compute selectivity; for filters this may be 1.0 on smalle samples;
+                # always put something slightly less than 1.0 to ensure that filters are pushed down when possible
                 selectivity = physical_op_df.passed_operator.sum() / num_source_records
+                op_name = physical_op_df.op_name.iloc[0].lower()
+                if selectivity == 1.0 and "filter" in op_name:
+                    selectivity -= 1e-3
 
                 # compute quality; if all qualities are None then this will be NaN
                 quality = physical_op_df.quality.mean()

palimpzest/query/processor/config.py

@@ -27,7 +27,7 @@ class QueryProcessorConfig(BaseModel):
     join_parallelism: int = Field(default=64)
     batch_size: int | None = Field(default=None)
     reasoning_effort: str | None = Field(default=None)  # Gemini: "disable", "low", "medium", "high"
-    use_vertex: bool = Field(default=True)  # Whether to use Vertex models for Gemini or Google models
+    use_vertex: bool = Field(default=False)  # Whether to use Vertex models for Gemini or Google models
     gemini_credentials_path: str | None = Field(default=None)  # Path to Gemini credentials file
     api_base: str | None = Field(default=None)  # API base URL for vLLM
 

palimpzest/query/processor/query_processor_factory.py

@@ -1,6 +1,9 @@
 import logging
+import os
 from enum import Enum
 
+from dotenv import load_dotenv
+
 from palimpzest.core.data.dataset import Dataset
 from palimpzest.core.elements.records import DataRecordCollection
 from palimpzest.query.execution.execution_strategy import ExecutionStrategy, SentinelExecutionStrategy
@@ -91,6 +94,27 @@ class QueryProcessorFactory:
         # set the final set of available models in the config
         config.available_models = available_models
 
+        if len(config.available_models) == 0:
+            raise ValueError("No available models found.")
+
+        openai_key = os.getenv("OPENAI_API_KEY")
+        anthropic_key = os.getenv("ANTHROPIC_API_KEY")
+        together_key = os.getenv("TOGETHER_API_KEY")
+        gemini_key = os.getenv("GEMINI_API_KEY")
+        google_key = os.getenv("GOOGLE_API_KEY")
+
+        for model in config.available_models:
+            if model.is_openai_model() and not openai_key:
+                raise ValueError("OPENAI_API_KEY must be set to use OpenAI models.")
+            if model.is_anthropic_model() and not anthropic_key:
+                raise ValueError("ANTHROPIC_API_KEY must be set to use Anthropic models.")
+            if model.is_together_model() and not together_key:
+                raise ValueError("TOGETHER_API_KEY must be set to use Together models.")
+            if model.is_google_ai_studio_model() and not (gemini_key or google_key or config.gemini_credentials_path):
+                raise ValueError("GEMINI_API_KEY, GOOGLE_API_KEY, or gemini_credentials path must be set to use Google Gemini models.")
+            if model.is_vllm_model() and config.api_base is None:
+                raise ValueError("api_base must be set to use vLLM models.")
+
         return config, validator
 
     @classmethod
@@ -172,6 +196,7 @@ class QueryProcessorFactory:
         train_dataset: dict[str, Dataset] | None = None,
         validator: Validator | None = None,
     ) -> DataRecordCollection:
+        load_dotenv(override=True)
         logger.info(f"Creating processor for dataset: {dataset}")
         processor = cls.create_processor(dataset, config, train_dataset, validator)
         logger.info(f"Created processor: {processor}")

palimpzest/utils/model_helpers.py

@@ -3,13 +3,12 @@ import os
 from palimpzest.constants import 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]:
+def get_models(include_embedding: bool = False, use_vertex: bool = False, 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.
     """
     models = []
-    if os.getenv("OPENAI_API_KEY") is not None:
+    if os.getenv("OPENAI_API_KEY") not in [None, ""]:
         openai_models = [model for model in Model if model.is_openai_model()]
         if not include_embedding:
             openai_models = [
@@ -17,7 +16,7 @@ def get_models(include_embedding: bool = False, use_vertex: bool = True, gemini_
             ]
         models.extend(openai_models)
 
-    if os.getenv("TOGETHER_API_KEY") is not None:
+    if os.getenv("TOGETHER_API_KEY") not in [None, ""]:
         together_models = [model for model in Model if model.is_together_model()]
         if not include_embedding:
             together_models = [
@@ -25,7 +24,7 @@ def get_models(include_embedding: bool = False, use_vertex: bool = True, gemini_
             ]
         models.extend(together_models)
 
-    if os.getenv("ANTHROPIC_API_KEY") is not None:
+    if os.getenv("ANTHROPIC_API_KEY") not in [None, ""]:
         anthropic_models = [model for model in Model if model.is_anthropic_model()]
         if not include_embedding:
             anthropic_models = [
@@ -38,9 +37,9 @@ def get_models(include_embedding: bool = False, use_vertex: bool = True, gemini_
         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):
+    if os.getenv("GEMINI_API_KEY") not in [None, ""] or (use_vertex and 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()]
+        google_ai_studio_models = [model for model in Model if model.is_google_ai_studio_model()]
         if not include_embedding:
             vertex_models = [
                 model for model in vertex_models if not model.is_embedding_model()
@@ -48,7 +47,7 @@ def get_models(include_embedding: bool = False, use_vertex: bool = True, gemini_
         if use_vertex:
             models.extend(vertex_models)
         else:
-            models.extend(google_models)
+            models.extend(google_ai_studio_models)
 
     if api_base is not None:
         vllm_models = [model for model in Model if model.is_vllm_model()]

{palimpzest-1.1.1.dist-info → palimpzest-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: palimpzest
-Version: 1.1.1
+Version: 1.3.0
 Summary: Palimpzest is a system which enables anyone to process AI-powered analytical queries simply by defining them in a declarative language
 Author-email: MIT DSG Semantic Management Lab <michjc@csail.mit.edu>
 Project-URL: homepage, https://palimpzest.org
@@ -31,9 +31,10 @@ Requires-Dist: pillow>=11.3.0
 Requires-Dist: prettytable>=3.9.0
 Requires-Dist: psutil==5.9.5
 Requires-Dist: PyLD>=2.0.4
-Requires-Dist: pyarrow==20.0.0
+Requires-Dist: pyarrow>=20.0.0
 Requires-Dist: pypdf>=5.1.0
 Requires-Dist: pytest-mock>=3.14.0
+Requires-Dist: python-dotenv>=1.2.1
 Requires-Dist: pyyaml>=6.0.1
 Requires-Dist: requests>=2.25
 Requires-Dist: ruff>=0.9.0

{palimpzest-1.1.1.dist-info → palimpzest-1.3.0.dist-info}/RECORD

@@ -1,11 +1,11 @@
 palimpzest/__init__.py,sha256=1PzadDDOVMQJKNEYUH0_tw8tQKUYTT31M0vuzTr2Rqk,1694
-palimpzest/constants.py,sha256=R_vNt3FW9MA4-HuszrGyGtDLRE9YMLXRMwX5GZcRH38,23323
+palimpzest/constants.py,sha256=kfuDTdSdeTCWp2JJopeuwpGIf8MC6jQ8SDBFuRIhN3k,23334
 palimpzest/policy.py,sha256=lIvw_C_rmwCH4LZaeNkAuixl8zw9RAW_JcSWSHPjKyc,11628
 palimpzest/agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/agents/compute_agents.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/agents/search_agents.py,sha256=t2QMreB5Ph71aoNk5bBtV-0l8im79z-pMAR3JDAySDw,29418
 palimpzest/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-palimpzest/core/models.py,sha256=FKyKW9PqmqpnDGWOINNT6XgBj0raaAskxtdNdFZ4Zyw,42688
+palimpzest/core/models.py,sha256=t4zHPA-Nrz2Mmq2EZfJWU_CsSbzu4LFv6_Wob10MZnc,46110
 palimpzest/core/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/core/data/context.py,sha256=x1xYyu9qW65dvtK_XayIfv_CgsCEPW6Qe0DTiSf9sjU,16207
 palimpzest/core/data/context_manager.py,sha256=8hAKWD2jhFZgghTu7AYgjkvKDsJUPVxq8g4nG0HWvfo,6150
@@ -38,7 +38,7 @@ palimpzest/query/execution/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 palimpzest/query/execution/all_sample_execution_strategy.py,sha256=d2MO_AmXF_HbV4rUNkFqHsuoBCofU98zQ-D3Q06BXCc,14369
 palimpzest/query/execution/execution_strategy.py,sha256=TnSInUlcGZHn2GUpLiIFVgfPpmsNfIKKgElnRt6a6ss,19248
 palimpzest/query/execution/execution_strategy_type.py,sha256=vRQBPCQN5_aoyD3TLIeW3VPo15mqF-5RBvEXkENz9FE,987
-palimpzest/query/execution/mab_execution_strategy.py,sha256=eKWEEFjOSpkaqkwuoFbg5yA5jdkki4FfRtCD7P4VaeI,48205
+palimpzest/query/execution/mab_execution_strategy.py,sha256=BLRTSQXPeWBlJ_-8GAFHj2fbIY_eoPhuWeDcIdOokcg,48247
 palimpzest/query/execution/parallel_execution_strategy.py,sha256=Di-8d7waE0bev4kNDXEJJqQ0wwQ87_sPV-t5qFtAlPQ,17589
 palimpzest/query/execution/single_threaded_execution_strategy.py,sha256=1rjMel0-AI6KUi_SMNgPPXxMgG5-t9lenLKoYEClgjk,17464
 palimpzest/query/generators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -46,23 +46,23 @@ palimpzest/query/generators/generators.py,sha256=OV2HBZvCQtvhj6rOwti_8dpQX_bqTZ-
 palimpzest/query/operators/__init__.py,sha256=B9zr_VmUs6YRep4fjbj7e0aTM6T9-GrqbY7tKWxEdkc,4734
 palimpzest/query/operators/aggregate.py,sha256=nQ6Zh1DYeqDoIDwkPQDw8QCwW0y52sGC-No6uPSRc8A,27367
 palimpzest/query/operators/compute.py,sha256=X_pWN45smg8L4dV54nOae7dldQGL1nJVlVyJ3ULWSmI,8432
-palimpzest/query/operators/convert.py,sha256=cjUPrSgvZBZXBbrbepIxZMBXjbWWPLuTX4JwLyvVg2U,16050
+palimpzest/query/operators/convert.py,sha256=beJLS-vnVc_VLnngoDKTj-k_Ul7GmDA-yMKM6-mX5Ho,16218
 palimpzest/query/operators/critique_and_refine.py,sha256=Q-NhasVoD9meX7g36RPrv3q4R48_8XEU4d3TE46hRJI,8979
 palimpzest/query/operators/distinct.py,sha256=ZTXlIS7IaFRTsWv9RemzCo1JLz25vEma-TB42CV5fJQ,2614
-palimpzest/query/operators/filter.py,sha256=ufREsO2-8CBk4u4fabDBYpEvb806E11EOyW-wuRs4vw,10356
-palimpzest/query/operators/join.py,sha256=CPqqIQrxWNjslWkxnZrbW-DKFkJ8iiDsWLJeZmUs5gI,38549
+palimpzest/query/operators/filter.py,sha256=h559CweLdcWw_-LPyR2h04pKsm-jVM_Kazif-BXYBFo,10516
+palimpzest/query/operators/join.py,sha256=hxIbSSDMy_5bch7kpYQ-iQS5XEth_T19cTSlXd76gdg,38845
 palimpzest/query/operators/limit.py,sha256=pdo7WfWY97SW3c-WqZ4SIPw7lHIVbaXPEWqHyK8qkF8,2130
 palimpzest/query/operators/logical.py,sha256=OtB82L1X19ibtLx1GIfeXXyO7YfjkFmh3puIUgqKQRE,21160
 palimpzest/query/operators/mixture_of_agents.py,sha256=KC-ZpjtGY28sfwlk2TpduLC_fADj_UASFCaicaKqSFc,11671
 palimpzest/query/operators/physical.py,sha256=0_BfFX9nKuN__440eAfEfApWAoGOunVSCZIQxS4HO2Y,9773
 palimpzest/query/operators/project.py,sha256=gxbjsHEACCP9wxATH-mw6wOSUi5s13JyWsLqqhAYWXQ,2111
-palimpzest/query/operators/rag.py,sha256=jJ09uCkjCr5KuiFMPhWcZAQNBR3_5eNb0ALkn_IoXbU,20157
+palimpzest/query/operators/rag.py,sha256=ZDloc4nC8foI3rTSHQxqduAVPj5LV8xMu_ng0EjDOA0,20409
 palimpzest/query/operators/scan.py,sha256=OqCiPRTvTY7SbauNMyFvGT5nRVeRzVsGYSrkoN1Ib_w,7407
 palimpzest/query/operators/search.py,sha256=cQin-Qc9FT7V0Gv3-pxMLbVMjqE6ALe99V0OrQhA6CI,22711
 palimpzest/query/operators/split.py,sha256=oLzwnYb8TNf3XA9TMKEAIw7EIA12wHneaD42BNLIHiI,15043
 palimpzest/query/operators/topk.py,sha256=MZl83Cu43QmN4skjlfpR8EVFFCgA7sR6PbGgBGWC0tg,13564
 palimpzest/query/optimizer/__init__.py,sha256=v9fSBOL2p3sQew4LrN2DQUPe0WezO328Hr54qBTqrAs,2799
-palimpzest/query/optimizer/cost_model.py,sha256=p7AsR6f4VYdGjrUKPGN_VTErY36GjY90Bsvsys4le2M,12655
+palimpzest/query/optimizer/cost_model.py,sha256=JaxdLuUZuq52BJ52YdW4ChfWptwXsh7Rk7oaPCn_gWc,12956
 palimpzest/query/optimizer/optimizer.py,sha256=ksLkzQ2sVgJFbkxGF3ncF74EsAHZFos8G19xlHQrtJo,20063
 palimpzest/query/optimizer/optimizer_strategy.py,sha256=0foDaBHqQehK_zz6IlDEbNIw-44wxY6LO5H1anJi56Y,10042
 palimpzest/query/optimizer/optimizer_strategy_type.py,sha256=V-MMHvJdnfZKoUX1xxxwh66q1RjN2FL35IsiT1C62c8,1084
@@ -71,9 +71,9 @@ palimpzest/query/optimizer/primitives.py,sha256=jMMVq37y1tWiPU1lSSKQP9OP-mzkpSxS
 palimpzest/query/optimizer/rules.py,sha256=awhe76trskv5Tq5E2QHpUN_YV6jH8INywa0Ige8IIhY,53341
 palimpzest/query/optimizer/tasks.py,sha256=DNJjY2QldfKFWj6INHElMh88dYc36Z5m3wHwbs4jyF4,30455
 palimpzest/query/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-palimpzest/query/processor/config.py,sha256=8-MpPYHv2SI4dub4MP_gOYSRxO80_ALLuWRxD-F2YOg,2521
+palimpzest/query/processor/config.py,sha256=rVpXNfzJvYstNZ2PxhQMWBpHCmjLDfiKh4ERKWRojc0,2522
 palimpzest/query/processor/query_processor.py,sha256=T4ffPbnOX23G8FDITzmM7Iw7DUEDWIHnwl8XLYllgjg,6240
-palimpzest/query/processor/query_processor_factory.py,sha256=i9L9StqlUi7m1AqZMuYQWhunqOJi3nLK47skhxq9tIA,8317
+palimpzest/query/processor/query_processor_factory.py,sha256=qED6pJtJJXVci2a4nScURSpJohrz1LIW6tXY05un2R8,9653
 palimpzest/schemabuilder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/schemabuilder/schema_builder.py,sha256=QraGp66dcD-ej6Y2mER40o86G9JqlBkL7swkJzjUAIY,7968
 palimpzest/tools/README.md,sha256=56_6LPG80uc0CLVhTBP6I1wgIffNv9cyTr0TmVZqmrM,483
@@ -84,13 +84,13 @@ palimpzest/tools/skema_tools.py,sha256=HXUFpjMhbVxZwKKkATeK-FwtlTCawaCbeP-uHntI1
 palimpzest/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/utils/env_helpers.py,sha256=n81KzoJ459pRxo7QmJA7duazwWsfoMGTHc71D2LatFk,334
 palimpzest/utils/hash_helpers.py,sha256=3A8dA7SbXTwnnvZvPVNqqMLlVRhCKyKF_bjNNAu3Exk,334
-palimpzest/utils/model_helpers.py,sha256=X6SlMgD5I5Aj_cxaFaoGaaNvOOqTNZVmjj6zbfn63Yk,2476
+palimpzest/utils/model_helpers.py,sha256=SqcY8rWzZ7D3Vgeq8d4OGNLGv4dXXVMWiJDqSCaalRQ,2490
 palimpzest/utils/progress.py,sha256=eHXrTPTCRHjMdK0EjYRUzSxcV6N1lK8TS3Ju_ZlQLhY,22002
 palimpzest/utils/udfs.py,sha256=LjHic54B1az-rKgNLur0wOpaz2ko_UodjLEJrazkxvY,1854
 palimpzest/validator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 palimpzest/validator/validator.py,sha256=SvjK09zCpGtK0yM0OasvQlSzyq3loy32DyOOKRmYXC0,15977
-palimpzest-1.1.1.dist-info/licenses/LICENSE,sha256=5GUlHy9lr-Py9kvV38FF1m3yy3NqM18fefuE9wkWumo,1079
-palimpzest-1.1.1.dist-info/METADATA,sha256=lZssP9vUbcnMNl_6vwpDvPyiI1IH65nDwggeSndSOcw,5359
-palimpzest-1.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-palimpzest-1.1.1.dist-info/top_level.txt,sha256=raV06dJUgohefUn3ZyJS2uqp_Y76EOLA9Y2e_fxt8Ew,11
-palimpzest-1.1.1.dist-info/RECORD,,
+palimpzest-1.3.0.dist-info/licenses/LICENSE,sha256=5GUlHy9lr-Py9kvV38FF1m3yy3NqM18fefuE9wkWumo,1079
+palimpzest-1.3.0.dist-info/METADATA,sha256=1VcCZTtJHg1ujb41-3QquXQLYaWtpCG9AHuvtrvJGvg,5395
+palimpzest-1.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+palimpzest-1.3.0.dist-info/top_level.txt,sha256=raV06dJUgohefUn3ZyJS2uqp_Y76EOLA9Y2e_fxt8Ew,11
+palimpzest-1.3.0.dist-info/RECORD,,