1  """
  2  Utility to enable or disable the JSONSimplifier post-processor
  3  ----------------------------------------------------------------------
  4  This script allows enabling or disabling the JSONSimplifier post-processor
  5  via the command line.
  6  """
  7  
  8  import os
  9  import sys
 10  import json
 11  import argparse
 12  from pathlib import Path
 13  
 14  def parse_args():
 15      """Parse command line arguments"""
 16      parser = argparse.ArgumentParser(
 17          description="Enable or disable the JSONSimplifier post-processor"
 18      )
 19      
 20      parser.add_argument(
 21          "--enable", 
 22          action="store_true",
 23          help="Enable the post-processor"
 24      )
 25      
 26      parser.add_argument(
 27          "--disable", 
 28          action="store_true",
 29          help="Disable the post-processor"
 30      )
 31      
 32      parser.add_argument(
 33          "--model",
 34          type=str,
 35          help="Specify the model to use"
 36      )
 37      
 38      parser.add_argument(
 39          "--apply-to",
 40          type=str,
 41          nargs="+",
 42          choices=["inference", "video", "transcription", "all"],
 43          help="Specify the task types to apply the post-processor to"
 44      )
 45      
 46      return parser.parse_args()
 47  
 48  def main():
 49      """Main function"""
 50      args = parse_args()
 51      
 52      # Check that either --enable or --disable is specified
 53      if not args.enable and not args.disable:
 54          print("Error: You must specify --enable or --disable")
 55          sys.exit(1)
 56      
 57      # Check that --enable and --disable are not both specified
 58      if args.enable and args.disable:
 59          print("Error: You cannot specify both --enable and --disable")
 60          sys.exit(1)
 61      
 62      # Load existing configuration
 63      config_path = os.environ.get("CONFIG_PATH", ".env")
 64      
 65      # Prepare modifications
 66      env_updates = {}
 67      
 68      # Update activation state
 69      if args.enable:
 70          env_updates["JSON_SIMPLIFIER_ENABLED"] = "true"
 71          print("Enabling JSONSimplifier post-processor")
 72      else:
 73          env_updates["JSON_SIMPLIFIER_ENABLED"] = "false"
 74          print("Disabling JSONSimplifier post-processor")
 75      
 76      # Update the model if specified
 77      if args.model:
 78          env_updates["JSON_SIMPLIFIER_MODEL"] = args.model
 79          print(f"Post-processor model set to: {args.model}")
 80      
 81      # Update task types if specified
 82      if args.apply_to:
 83          if "all" in args.apply_to:
 84              tasks = ["inference", "video", "transcription"]
 85          else:
 86              tasks = args.apply_to
 87          
 88          env_updates["JSON_SIMPLIFIER_APPLY_TO"] = ",".join(tasks)
 89          print(f"The post-processor will be applied to task types: {', '.join(tasks)}")
 90      
 91      # Update the .env file
 92      try:
 93          # Load existing content
 94          env_content = {}
 95          if Path(config_path).exists():
 96              with open(config_path, "r") as f:
 97                  for line in f:
 98                      line = line.strip()
 99                      if line and not line.startswith("#") and "=" in line:
100                          key, value = line.split("=", 1)
101                          env_content[key.strip()] = value.strip()
102          
103          # Update with new values
104          env_content.update(env_updates)
105          
106          # Write updated file
107          with open(config_path, "w") as f:
108              for key, value in env_content.items():
109                  f.write(f"{key}={value}\n")
110          
111          print(f"Configuration updated in {config_path}")
112          print("Restart the API to apply changes")
113          
114      except Exception as e:
115          print(f"Error updating configuration: {str(e)}")
116          sys.exit(1)
117  
118  if __name__ == "__main__":
119      main()