content: """Implements `Check` objects used for LLM evaluation purposes.
  """
  
  from __future__ import annotations
  from dataclasses import dataclass
  import os
  import typing as t
  
  from loguru import logger
  import polars as pl
  from pydantic import BaseModel
  
  from uptrain.operators.base import (
      Operator,
      TransformOp,
      ColumnOp,
      deserialize_operator,
  )
  from uptrain.utilities import jsonload, jsondump, to_py_types, clear_directory
  from uptrain.framework.base import OperatorDAG, Settings
  
  __all__ = ["Check", "CheckSet", "ExperimentArgs"]
  
  
  class Check:
      """A simple check that runs the given list of table operators in sequence.
  
      Attributes:
          name (str): Name of the check.
          operators (list[TableOp]): A list of operators to run in sequence on the input data. The output of each
              operator is passed as input to the next operator.
          plots (list[Operator]): How to plot the output of the check.
      """
  
      name: str
      operators: list[Operator]
      plots: list[Operator]
  
      def __init__(
          self,
          name: str,
          operators: list[Operator],
          plots: t.Union[list[Operator], None] = None,
      ):
          self.name = name
          self.operators = operators
          self.plots = plots if plots is not None else []
  
      def setup(self, settings: "Settings"):
          self._settings = settings
  
          # no need to add the plot operator to the dag, since it's run later
          self._op_dag = OperatorDAG(name=self.name)
          for i, op in enumerate(self.operators):
              if i == 0:
                  deps = []
              else:
                  deps = [f"operator_{i-1}"]
              self._op_dag.add_step(f"operator_{i}", op, deps=deps)
          self._op_dag.setup(settings)
  
          return self
  
      def run(
          self, data: t.Union[pl.DataFrame, None] = None
      ) -> t.Union[pl.DataFrame, None]:
          """Run this check on the given data."""
          node_inputs = {"operator_0": data}
  
          if len(self.operators):
              # pick output from the last op in the sequence
              name_final_node = f"operator_{len(self.operators) - 1}"
              node_outputs = self._op_dag.run(
                  node_inputs=node_inputs,
                  output_nodes=[name_final_node],
              )
              return node_outputs[name_final_node]
          else:
              return data
  
      def dict(self) -> dict:
          """Serialize this check to a dict."""
          return {
              "name": self.name,
              "operators": [to_py_types(op) for op in self.operators],
              "plots": [to_py_types(op) for op in self.plots],
          }  # serializes only the attributes of the class, like pydantic models
  
      @classmethod
      def from_dict(cls, data: dict) -> "Check":
          """Deserialize a check from a dict of its parameters."""
          operators = [deserialize_operator(op) for op in data["operators"]]
          plots = [deserialize_operator(op) for op in data["plots"]]
          return cls(name=data["name"], operators=operators, plots=plots)  # type: ignore
  
  
  class CheckSet:
      """Container for a set of checks to run together. This is the entrypoint to Uptrain for users.
  
      Attributes:
          source (Operator): The source operator to run. Specifies where to get the data from.
          preprocessors (list[TableOp]): A list of operators to run on the input data before running the checks.
          checks (list[Check]): The set of checks to run on the input data.
      """
  
      source: Operator
      checks: list[Check]
      preprocessors: list[TransformOp]
      postprocessors: list[Operator]
  
      def __init__(
          self,
          source: Operator,
          checks: list[t.Any],
          preprocessors: t.Union[list[TransformOp], None] = None,
          postprocessors: t.Union[list[Operator], None] = None,
      ):
          self.source = source
          self.checks = checks
          self.preprocessors = preprocessors if preprocessors is not None else []
          self.postprocessors = postprocessors if postprocessors is not None else []
  
          # verify all checks have different names
          check_names = [check.name for check in checks]
          assert len(set(check_names)) == len(check_names), "Duplicate check names"
          for check in checks:
              assert isinstance(check, Check), "Each check must be an instance of Check"
  
      def setup(self, settings: Settings):
          """Create the logs directory, or clear it if it already exists. Also, persist the
          evaluation config.
          """
          self._settings = settings
          logs_dir = self._settings.logs_folder
          if not os.path.exists(logs_dir):
              os.makedirs(logs_dir)
          else:
              clear_directory(logs_dir)
  
          logger.info(f"Uptrain Logs directory: {logs_dir}")
  
          # persist the check-set as well as the corresponding settings
          self.serialize(os.path.join(logs_dir, "config.json"))
          self._settings.serialize(os.path.join(logs_dir, "settings.json"))
  
          self.source.setup(self._settings)
          for preprocessor in self.preprocessors:
              preprocessor.setup(self._settings)
          for check in self.checks:
              check.setup(self._settings)
          for postprocessor in self.postprocessors:
              postprocessor.setup(self._settings)
          return self
  
      def run(self):
          """Run all checks in this set."""
          from uptrain.operators import JsonWriter
  
          logger.info("CheckSet Status: Starting checkset")
  
          source_output = self.source.run()["output"]
          if source_output is None:
              raise RuntimeError("Dataset read from the source is: None")
          if len(source_output) == 0:
              raise RuntimeError("Dataset read from the source is: empty")
          logger.info("CheckSet Status: Dataset loaded from source")
  
          if len(self.preprocessors) > 0:
              for preprocessor in self.preprocessors:
                  source_output = preprocessor.run(source_output)["output"]
                  assert source_output is not None, "Output of preprocessor is None"
  
              # persist the preprocessed input for debugging
              JsonWriter(
                  fpath=os.path.join(
                      self._settings.logs_folder, "preprocessed_input.jsonl"
                  )
              ).setup(self._settings).run(source_output)
  
          logger.info("CheckSet Status: Preprocessing Done")
  
          consolidated_output = {}
          for check in self.checks:
              logger.info(f"CheckSet Status: Check {check.name} Started")
              check_output = check.run(source_output)
              assert check_output is not None, f"Output of check {check.name} is None"
              self._get_sink_for_check(self._settings, check).run(check_output)
              logger.info(f"CheckSet Status: Check {check.name} Completed")
  
              if len(self.postprocessors):
                  if not all(isinstance(op, ColumnOp) for op in check.operators):
                      continue
                  for col in check_output.columns:
                      consolidated_output[col] = check_output[col]
          logger.info("CheckSet Status: All Checks Completed")
  
          if len(self.postprocessors):
              consolidated_output = pl.DataFrame(consolidated_output)
              for postprocessor in self.postprocessors:
                  consolidated_output = postprocessor.run(consolidated_output)["output"]
                  assert (
                      consolidated_output is not None
                  ), "Output of postprocessor is None"
  
              # persist the postprocessed input for debugging
              JsonWriter(
                  fpath=os.path.join(
                      self._settings.logs_folder, "postprocessed_input.jsonl"
                  )
              ).setup(self._settings).run(consolidated_output)
          logger.info("CheckSet Status: Postprocessing Done")
  
      @staticmethod
      def _get_sink_for_check(settings: Settings, check: Check):
          """Get the sink operator for this check."""
          from uptrain.operators import JsonWriter
  
          fname = check.name.replace(" ", "_") + ".jsonl"
          return JsonWriter(fpath=os.path.join(settings.logs_folder, fname))
  
      @classmethod
      def from_dict(cls, data: dict) -> "CheckSet":
          checks = [Check.from_dict(check) for check in data.get("checks", [])]
          return cls(
              source=deserialize_operator(data["source"]),
              preprocessors=[
                  deserialize_operator(op) for op in data.get("preprocessors", [])
              ],  # type: ignore
              checks=checks,
              postprocessors=[
                  deserialize_operator(op) for op in data.get("postprocessors", [])
              ],  # type: ignore
          )
  
      def dict(self) -> dict:
          return {
              "source": to_py_types(self.source),
              "preprocessors": [to_py_types(op) for op in self.preprocessors],
              "checks": [check.dict() for check in self.checks],
              "postprocessors": [to_py_types(op) for op in self.postprocessors],
          }
  
      @classmethod
      def deserialize(cls, fpath: str) -> "CheckSet":
          with open(fpath, "r") as f:
              return cls.from_dict(jsonload(f))
  
      def serialize(self, fpath: t.Optional[str] = None):
          """Serialize this check set along with the run settings to a JSON file."""
          if fpath is None:
              fpath = os.path.join(self._settings.logs_folder, "config.json")
  
          with open(fpath, "w") as f:
              jsondump(self.dict(), f)
  
  
  class ExperimentArgs(BaseModel):
      """Container for arguments to run a LLM experiment. This is the entrypoint to run prompt/model tests with Uptrain for users.
      This generates multiple prompts, generate LLM responses for them and parses the response to add multiple columns.
  
      Attributes:
          prompt_template (str): Prompt template to generate multiple prompts to experiment with.
          prompt_params (dict(str, str)): Dictionary of prompt variables with experiment with.
          models (str or list[str]): Either a single model or list of models to experiment with.
          context_vars (dict or list): List of dataset variables with mapping between variable name in prompt and variable name in dataset
          temperature (float): Temperature for the LLM response generation.
          col_out_response (str): LLM output will be saved under this column name
          col_out_mapping (dict): Mapping used to parse LLM response
      """
  
      prompt_template: str
      prompt_params: dict = {}
      models: t.Union[str, list[str]]
      context_vars: t.Union[dict, list]
      temperature: float = 1.0
      col_out_response: str = "response"
      col_out_mapping: t.Optional[dict] = None
  
      def _get_preprocessors(self):
          from uptrain.operators import PromptGenerator, TextCompletion, OutputParser
  
          "Convert experiment args into checkset preprocessors for execution"
          if isinstance(self.context_vars, list):
              self.context_vars = dict(zip(self.context_vars, self.context_vars))
  
          preprocessors = [
              PromptGenerator(
                  prompt_template=self.prompt_template,
                  prompt_params=self.prompt_params,
                  models=[self.models] if isinstance(self.models, str) else self.models,
                  context_vars=self.context_vars,
              ),
              TextCompletion(
                  col_in_prompt="exp_prompt",
                  col_in_model="exp_model",
                  col_out_completion=self.col_out_response,
                  temperature=self.temperature,
              ),
          ]
  
          if self.col_out_mapping is not None:
              preprocessors.append(
                  OutputParser(
                      col_in_response=self.col_out_response,
                      col_out_mapping=self.col_out_mapping,
                  )
              )
  
          return preprocessors
  
      def _modify_checks(self, checks):
          "Modify checks to add experiment variables into plots"
          for check in checks:
              if check.plots is not None:
                  for plot in check.plots:
                      if plot.kind in ["histogram", "bar", "line"]:
                          if "barmode" not in plot.props:
                              plot.props.update({"barmode": "group"})
                          plot.props["color"] = "exp_experiment_id"
          return checks
  