palimpzest 1.1.0__tar.gz → 1.2.0__tar.gz

{palimpzest-1.1.0/src/palimpzest.egg-info → palimpzest-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: palimpzest
-Version: 1.1.0
+Version: 1.2.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,7 +31,7 @@ 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: pyyaml>=6.0.1

{palimpzest-1.1.0 → palimpzest-1.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "palimpzest"
-version = "1.1.0"
+version = "1.2.0"
 description = "Palimpzest is a system which enables anyone to process AI-powered analytical queries simply by defining them in a declarative language"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -25,7 +25,7 @@ dependencies = [
     "prettytable>=3.9.0",
     "psutil==5.9.5",
     "PyLD>=2.0.4",
-    "pyarrow==20.0.0",
+    "pyarrow>=20.0.0",
     "pypdf>=5.1.0",
     "pytest-mock>=3.14.0",
     "pyyaml>=6.0.1",

{palimpzest-1.1.0 → palimpzest-1.2.0}/src/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-1.1.0 → palimpzest-1.2.0}/src/palimpzest/prompts/prompt_factory.py

@@ -830,7 +830,7 @@ class PromptFactory:
                 field_type = dr.get_field_type(field_name)
 
                 # audio filepath (or list of audio filepaths)
-                if field_type.annotation in [AudioFilepath, AudioFilepath | None, AudioFilepath | Any]:
+                if field_type.annotation in [AudioFilepath, AudioFilepath | None, AudioFilepath | Any] and field_value is not None:
                     with open(field_value, "rb") as f:
                         base64_audio_str = base64.b64encode(f.read()).decode("utf-8")
                     audio_content.append(
@@ -839,6 +839,8 @@ class PromptFactory:
 
                 elif field_type.annotation in [list[AudioFilepath], list[AudioFilepath] | None, list[AudioFilepath] | Any]:
                     for audio_filepath in field_value:
+                        if audio_filepath is None:
+                            continue
                         with open(audio_filepath, "rb") as f:
                             base64_audio_str = base64.b64encode(f.read()).decode("utf-8")
                         audio_content.append(
@@ -846,13 +848,15 @@ class PromptFactory:
                         )
 
                 # pre-encoded images (or list of pre-encoded images)
-                elif field_type.annotation in [AudioBase64, AudioBase64 | None, AudioBase64 | Any]:
+                elif field_type.annotation in [AudioBase64, AudioBase64 | None, AudioBase64 | Any] and field_value is not None:
                     audio_content.append(
                         {"type": "input_audio", "input_audio": {"data": field_value, "format": "wav"}}
                     )
 
                 elif field_type.annotation in [list[AudioBase64], list[AudioBase64] | None, list[AudioBase64] | Any]:
                     for base64_audio in field_value:
+                        if base64_audio is None:
+                            continue
                         audio_content.append(
                             {"type": "input_audio", "input_audio": {"data": base64_audio, "format": "wav"}}
                         )
@@ -882,7 +886,7 @@ class PromptFactory:
                 field_type = dr.get_field_type(field_name)
 
                 # image filepath (or list of image filepaths)
-                if field_type.annotation in [ImageFilepath, ImageFilepath | None, ImageFilepath | Any]:
+                if field_type.annotation in [ImageFilepath, ImageFilepath | None, ImageFilepath | Any] and field_value is not None:
                     with open(field_value, "rb") as f:
                         base64_image_str = base64.b64encode(f.read()).decode("utf-8")
                     image_content.append(
@@ -891,6 +895,8 @@ class PromptFactory:
 
                 elif field_type.annotation in [list[ImageFilepath], list[ImageFilepath] | None, list[ImageFilepath] | Any]:
                     for image_filepath in field_value:
+                        if image_filepath is None:
+                            continue
                         with open(image_filepath, "rb") as f:
                             base64_image_str = base64.b64encode(f.read()).decode("utf-8")
                         image_content.append(
@@ -898,21 +904,25 @@ class PromptFactory:
                         )
 
                 # image url (or list of image urls)
-                elif field_type.annotation in [ImageURL, ImageURL | None, ImageURL | Any]:
+                elif field_type.annotation in [ImageURL, ImageURL | None, ImageURL | Any] and field_value is not None:
                     image_content.append({"type": "image_url", "image_url": {"url": field_value}})
 
                 elif field_type.annotation in [list[ImageURL], list[ImageURL] | None, list[ImageURL] | Any]:
                     for image_url in field_value:
+                        if image_url is None:
+                            continue
                         image_content.append({"type": "image_url", "image_url": {"url": image_url}})
 
                 # pre-encoded images (or list of pre-encoded images)
-                elif field_type.annotation in [ImageBase64, ImageBase64 | None, ImageBase64 | Any]:
+                elif field_type.annotation in [ImageBase64, ImageBase64 | None, ImageBase64 | Any] and field_value is not None:
                     image_content.append(
                         {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{field_value}"}}
                     )
 
                 elif field_type.annotation in [list[ImageBase64], list[ImageBase64] | None, list[ImageBase64] | Any]:
                     for base64_image in field_value:
+                        if base64_image is None:
+                            continue
                         image_content.append(
                             {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                         )

{palimpzest-1.1.0 → palimpzest-1.2.0}/src/palimpzest/query/execution/execution_strategy.py

@@ -91,6 +91,7 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
         use_final_op_quality: bool = False,
         seed: int = 42,
         exp_name: str | None = None,
+        dont_use_priors: bool = False,
         *args,
         **kwargs,
     ):
@@ -105,6 +106,7 @@ class SentinelExecutionStrategy(BaseExecutionStrategy, ABC):
         self.seed = seed
         self.rng = np.random.default_rng(seed=seed)
         self.exp_name = exp_name
+        self.dont_use_priors = dont_use_priors
 
         # general cache which maps hash(logical_op_id, phys_op_id, hash(input)) --> record_set
         self.cache: dict[int, DataRecordSet] = {}

{palimpzest-1.1.0 → palimpzest-1.2.0}/src/palimpzest/query/execution/mab_execution_strategy.py

@@ -44,6 +44,7 @@ class OpFrontier:
             seed: int,
             policy: Policy,
             priors: dict | None = None,
+            dont_use_priors: bool = False,
         ):
         # set k and j, which are the initial number of operators in the frontier and the
         # initial number of records to sample for each frontier operator
@@ -51,6 +52,7 @@ class OpFrontier:
         self.j = j
         self.source_indices = source_indices
         self.root_dataset_ids = root_dataset_ids
+        self.dont_use_priors = dont_use_priors
 
         # store the policy that we are optimizing under
         self.policy = policy
@@ -68,6 +70,7 @@ class OpFrontier:
         is_llm_filter = isinstance(sample_op, LLMFilter)
         is_llm_topk = isinstance(sample_op, TopKOp) and isinstance(sample_op.index, Collection)
         self.is_llm_op = is_llm_convert or is_llm_filter or is_llm_topk or self.is_llm_join
+        self.is_llm_convert = is_llm_convert
 
         # get order in which we will sample physical operators for this logical operator
         sample_op_indices = self._get_op_index_order(op_set, seed)
@@ -190,7 +193,9 @@ class OpFrontier:
         Returns a list of indices for the operators in the op_set.
         """
         # if this is not an llm-operator, we simply return the indices in random order
-        if not self.is_llm_op:
+        if not self.is_llm_op or self.dont_use_priors:
+            if self.is_llm_convert:
+                print("Using NO PRIORS for operator sampling order")
             rng = np.random.default_rng(seed=seed)
             op_indices = np.arange(len(op_set))
             rng.shuffle(op_indices)
@@ -198,6 +203,8 @@ class OpFrontier:
 
         # if this is an llm-operator, but we do not have priors, we first compute naive priors
         if self.priors is None or any([op_id not in self.priors for op_id in map(lambda op: op.get_op_id(), op_set)]):
+            if self.is_llm_convert:
+                print("Using NAIVE PRIORS for operator sampling order")
             self.priors = self._compute_naive_priors(op_set)
 
         # NOTE: self.priors is a dictionary with format:
@@ -770,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
@@ -805,7 +812,7 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
                 assert len(root_dataset_ids) == 1, f"Scan for {sample_op} has {len(root_dataset_ids)} > 1 root dataset ids"
                 root_dataset_id = root_dataset_ids[0]
                 source_indices = dataset_id_to_shuffled_source_indices[root_dataset_id]
-                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors)
+                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors, self.dont_use_priors)
             elif isinstance(sample_op, JoinOp):
                 assert len(source_unique_logical_op_ids) == 2, f"Join for {sample_op} has {len(source_unique_logical_op_ids)} != 2 source logical operators"
                 left_source_indices = op_frontiers[source_unique_logical_op_ids[0]].source_indices
@@ -814,10 +821,10 @@ class MABExecutionStrategy(SentinelExecutionStrategy):
                 for left_source_idx in left_source_indices:
                     for right_source_idx in right_source_indices:
                         source_indices.append((left_source_idx, right_source_idx))
-                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors)
+                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors, self.dont_use_priors)
             else:
                 source_indices = op_frontiers[source_unique_logical_op_ids[0]].source_indices
-                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors)
+                op_frontiers[unique_logical_op_id] = OpFrontier(op_set, source_unique_logical_op_ids, root_dataset_ids, source_indices, self.k, self.j, self.seed, self.policy, self.priors, self.dont_use_priors)
 
         # initialize and start the progress manager
         self.progress_manager = create_progress_manager(plan, sample_budget=self.sample_budget, sample_cost_budget=self.sample_cost_budget, progress=self.progress)

{palimpzest-1.1.0 → palimpzest-1.2.0}/src/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-1.1.0 → palimpzest-1.2.0}/src/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,