feat: --format and --force flag for CLI (#1310)

## Summary

For the VS Code plugin we are currently doing parsing and cleaning on
text in order to understand what to prompt to user or what the status is
post execution (errors, etc).

Adding two things to simplify life:
--force flag -> forcefully compiles ignoring prompt... aka demo mode.
--format [JSON| TEXT] --> Only output a json that can be parsed and
operated on rather than user friendly text.

To achieve this:
* added a format_print(... ,format=...) method to wrap our prints and
skip them when json mode.
* added a format enum, decorator to wrap exceptions.
* added format mostly everywhere.


TODO: Still need to cover more test of the cli to make sure it works as
intended.


## Checklist
- [ ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added TEXT/JSON CLI output modes with unified formatter, structured
JSON errors/prompts, and a --force option; formatting applied across
compile, hub, and uploader flows.

* **Tests**
* Added a JSON-format compile test and an autouse fixture to improve
test import/path isolation.

* **Chores**
* Switched many test/sample artifacts to sample_sources, added numerous
compiled sample/canary data files, and minor build/test config updates.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: cristian-zlai

python/README.md

@@ -13,7 +13,7 @@ Chronon Python API for materializing configs to be run by the Chronon Engine. Co
 Most fields are self explanatory. Time columns are expected to be in milliseconds (unixtime).
 
 ```python
-# File <repo>/sources/test_sources.py
+# File <repo>/sources/sample_sources.py
 from ai.chronon.query import (
   Query,
   select,
@@ -65,7 +65,7 @@ from ai.chronon.group_by import (
   Aggregation,
   DefaultAggregation,
 )
-from sources import test_sources
+from sources import sample_sources
 
 sum_cols = [f"active_{x}_days" for x in [30, 90, 120]]
 
@@ -114,11 +114,11 @@ A Join is a collection of feature values for the keys and (times if applicable)
 ```python
 # File <repo>/joins/example_team/example_join.py
 from ai.chronon.join import Join, JoinPart
-from sources import test_sources
+from sources import sample_sources
 from group_bys.example_team import example_group_by
 
 v1 = Join(
-    left=test_sources.website__views,
+    left=sample_sources.website__views,
     right_parts=[
         JoinPart(group_by=example_group_by.v0),
     ],

python/package.mill

@@ -63,15 +63,14 @@ object `package` extends PythonModule with RuffModule with PublishModule {
 
     override def sources = Task.Sources(moduleDir)
 
-    // Override forkEnv to set PYTHONPATH
+    // Override forkEnv to set PYTHONPATH for test isolation
     override def forkEnv: T[Map[String, String]] = Task {
 
       val generatedPath = build.python.generatedSources().head.path.toString
       val sourcePath = (moduleDir / os.up / "src").toString
-      val samplePath = (moduleDir / "sample").toString
-
+      // Note: samplePath removed from global PYTHONPATH to avoid conflicts with canary/
       super.forkEnv() ++ Map(
-        "PYTHONPATH" -> s"$generatedPath:$sourcePath:$samplePath"
+        "PYTHONPATH" -> s"$generatedPath:$sourcePath",
       )
 
     }

python/src/ai/chronon/cli/compile/compile_context.py

@@ -17,6 +17,7 @@
 from ai.chronon.cli.compile.display.compile_status import CompileStatus
 from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
 from ai.chronon.cli.compile.serializer import file2thrift
+from ai.chronon.cli.formatter import Format
 from ai.chronon.cli.logger import get_logger, require
 
 logger = get_logger()
@@ -31,11 +32,13 @@ class ConfigInfo:
 
 @dataclass
 class CompileContext:
-    def __init__(self, ignore_python_errors: bool = False):
+    def __init__(self, ignore_python_errors: bool = False, format: Format = Format.TEXT, force: bool = False):
         self.chronon_root: str = os.getenv("CHRONON_ROOT", os.getcwd())
-        self.teams_dict: Dict[str, Team] = teams.load_teams(self.chronon_root)
+        self.teams_dict: Dict[str, Team] = teams.load_teams(self.chronon_root, print=format != Format.JSON)
         self.compile_dir: str = "compiled"
         self.ignore_python_errors: bool = ignore_python_errors
+        self.format: Format = format
+        self.force: bool = force
 
         self.config_infos: List[ConfigInfo] = [
             ConfigInfo(folder_name="joins", cls=Join, config_type=ConfType.JOIN),
@@ -55,7 +58,7 @@ def __init__(self, ignore_python_errors: bool = False):
             ),  # only for team metadata
         ]
 
-        self.compile_status = CompileStatus(use_live=False)
+        self.compile_status = CompileStatus(use_live=False, format=format)
 
         self.existing_confs: Dict[Type, Dict[str, Any]] = {}
         for config_info in self.config_infos:

python/src/ai/chronon/cli/compile/compiler.py

@@ -14,6 +14,7 @@
 from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
 from ai.chronon.cli.compile.display.console import console
 from ai.chronon.cli.compile.parse_teams import merge_team_execution_info
+from ai.chronon.cli.formatter import Format, PromptException
 from ai.chronon.types import MetaData
 
 logger = logger.get_logger()
@@ -52,14 +53,26 @@ def compile(self) -> Dict[ConfType, CompileResult]:
         self.compile_context.validator.validate_changes(all_compiled_objects)
 
         # Show the nice display first
-        console.print(
-            self.compile_context.compile_status.render(self.compile_context.ignore_python_errors)
-        )
+        if self.compile_context.format != Format.JSON:
+            console.print(
+                self.compile_context.compile_status.render(self.compile_context.ignore_python_errors)
+            )
 
         # Check for confirmation before finalizing files
-        self.compile_context.validator.check_pending_changes_confirmation(
-            self.compile_context.compile_status
-        )
+        if not self.compile_context.force:
+            if self.compile_context.format != Format.JSON:
+                self.compile_context.validator.check_pending_changes_confirmation(
+                    self.compile_context.compile_status
+                )
+            else:
+                # In case of JSON format we need to prompt the user for confirmation if changes are not versioned.
+                non_version_changes = self.compile_context.validator._non_version_changes()
+                if non_version_changes:
+                    raise PromptException(prompt=f"The following configs are changing in-place (changing semantics without changing the version)."
+                    f" {', '.join([v.name for v in non_version_changes])} Do you want to proceed?", 
+                    options= ["yes", "no"],
+                    instructions="If 'yes' run with --force to proceed with the compilation."
+                    )
 
         # Only proceed with file operations if there are no compilation errors
         if not self._has_compilation_errors() or self.compile_context.ignore_python_errors:
@@ -151,23 +164,25 @@ def _write_objects_in_folder(
                 if co.errors:
                     error_dict[co.name] = co.errors
 
-                    for error in co.errors:
-                        self.compile_context.compile_status.print_live_console(
-                            f"Error processing conf {co.name}: {error}"
-                        )
-                        traceback.print_exception(type(error), error, error.__traceback__)
+                    if self.compile_context.format != Format.JSON:
+                        for error in co.errors:
+                            self.compile_context.compile_status.print_live_console(
+                                f"Error processing conf {co.name}: {error}"
+                            )
+                            traceback.print_exception(type(error), error, error.__traceback__)
 
                 else:
                     self._write_object(co)
                     object_dict[co.name] = co.obj
             else:
                 error_dict[co.file] = co.errors
 
-                self.compile_context.compile_status.print_live_console(
-                    f"Error processing file {co.file}: {co.errors}"
-                )
-                for error in co.errors:
-                    traceback.print_exception(type(error), error, error.__traceback__)
+                if self.compile_context.format != Format.JSON:
+                    self.compile_context.compile_status.print_live_console(
+                        f"Error processing file {co.file}: {co.errors}"
+                    )
+                    for error in co.errors:
+                        traceback.print_exception(type(error), error, error.__traceback__)
 
         return object_dict, error_dict
 

python/src/ai/chronon/cli/compile/conf_validator.py

@@ -672,16 +672,18 @@ def check_pending_changes_confirmation(self, compile_status):
             return  # No pending changes
 
         # Check if we need user confirmation (only for non-version-bump changes)
-        non_version_changes = self._filter_non_version_changes(
-            self._pending_changes["changed"] + self._pending_changes["deleted"],
-            self._pending_changes["added"],
-        )
-
-        if non_version_changes:
+        if self._non_version_changes():
             if not self._prompt_user_confirmation():
                 console.print("❌ Compilation cancelled by user.")
                 sys.exit(1)
 
+    def _non_version_changes(self):
+        """Return list of changes that are NOT version bumps and require confirmation."""
+        return self._filter_non_version_changes(
+            self._pending_changes["changed"] + self._pending_changes["deleted"],
+            self._pending_changes["added"],
+        )
+
     def _has_compilation_errors(self, compile_status):
         """Check if there are any compilation errors across all class trackers."""
         for tracker in compile_status.cls_to_tracker.values():

python/src/ai/chronon/cli/compile/display/class_tracker.py

@@ -5,20 +5,22 @@
 
 from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
 from ai.chronon.cli.compile.display.diff_result import DiffResult
+from ai.chronon.cli.formatter import Format, format_print
 
 
 class ClassTracker:
     """
     Tracker object per class - Join, StagingQuery, GroupBy etc
     """
 
-    def __init__(self):
+    def __init__(self, format: Format = Format.TEXT):
         self.existing_objs: Dict[str, CompiledObj] = {}  # name to obj
         self.files_to_obj: Dict[str, List[Any]] = {}
         self.files_to_errors: Dict[str, List[Exception]] = {}
         self.new_objs: Dict[str, CompiledObj] = {}  # name to obj
         self.diff_result = DiffResult()
         self.deleted_names: List[str] = []
+        self.format = format
 
     def add_existing(self, obj: CompiledObj) -> None:
         self.existing_objs[obj.name] = obj
@@ -51,9 +53,9 @@ def _update_diff(self, compiled: CompiledObj) -> None:
                     n=2,
                 )
 
-                print(f"Updated object: {compiled.name} in file {compiled.file}")
-                print("".join(diff))
-                print("\n")
+                format_print(f"Updated object: {compiled.name} in file {compiled.file}", format=self.format)
+                format_print("".join(diff), format=self.format)
+                format_print("\n", format=self.format)
 
                 self.diff_result.updated.append(compiled.name)
 

python/src/ai/chronon/cli/compile/display/compile_status.py

@@ -6,21 +6,23 @@
 
 from ai.chronon.cli.compile.display.class_tracker import ClassTracker
 from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.formatter import Format
 
 
 class CompileStatus:
     """
     Uses rich ui - to consolidate and sink the overview of the compile process to the bottom.
     """
 
-    def __init__(self, use_live: bool = False):
+    def __init__(self, use_live: bool = False, format: Format = Format.TEXT):
         self.cls_to_tracker: Dict[str, ClassTracker] = OrderedDict()
         self.use_live = use_live
         # we need vertical_overflow to be visible as the output gets cufoff when our output goes past the termianal window
         # but then we start seeing duplicates: https://github.com/Textualize/rich/issues/3263
         if self.use_live:
             self.live = Live(refresh_per_second=50, vertical_overflow="visible")
             self.live.start()
+        self.format = format
 
     def print_live_console(self, msg: str):
         if self.use_live:
@@ -33,7 +35,7 @@ def add_object_update_display(self, compiled: CompiledObj, obj_type: str = None)
             )
 
         if obj_type not in self.cls_to_tracker:
-            self.cls_to_tracker[obj_type] = ClassTracker()
+            self.cls_to_tracker[obj_type] = ClassTracker(format=self.format)
 
         self.cls_to_tracker[obj_type].add(compiled)
 
@@ -43,7 +45,7 @@ def add_existing_object_update_display(self, existing_obj: CompiledObj) -> None:
         obj_type = existing_obj.obj_type
 
         if obj_type not in self.cls_to_tracker:
-            self.cls_to_tracker[obj_type] = ClassTracker()
+            self.cls_to_tracker[obj_type] = ClassTracker(format=self.format)
 
         self.cls_to_tracker[obj_type].add_existing(existing_obj)
 

python/src/ai/chronon/cli/formatter.py

@@ -0,0 +1,65 @@
+"""
+Formatter for CLI responses.
+
+In the interest of making our CLI MCP friendly as well as VSCode plugin friendly, we need more structured responses to avoid parsing the output manually.
+
+For this purpose we add a new kind of exception that can be raised to prompt the user for a confirmation.
+Any exception is caught by the jsonify_exceptions_if_json_format decorator and turned into a JSON response.
+"""
+
+import functools
+import json
+import sys
+import traceback
+from enum import Enum
+
+
+class PromptException(Exception):
+    """
+    Exception to raise when user prompt is needed.
+    Args:
+        prompt: The prompt to display to the user.
+        options: The options to display to the user.
+        instructions: The instructions for implementation.
+    """
+    def __init__(self, prompt, options, instructions):
+        super().__init__(prompt)
+        self.prompt = prompt
+        self.options = options
+        self.instructions = instructions
+
+class Format(Enum):
+    JSON = "json"
+    TEXT = "text"
+
+def format_print(lines, format: Format = Format.TEXT):
+    """ Format aware print. """
+    if format != Format.JSON:
+        print(lines)
+
+
+def jsonify_exceptions_if_json_format(func):
+    """ Turn exceptions into JSON if the format is JSON. If a prompt exception is raised, return the prompt and options."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        use_json = kwargs.get("format") == Format.JSON
+        try:
+            return func(*args, **kwargs)
+        except PromptException as e:
+            if use_json:
+                print(json.dumps({
+                    "prompt": e.prompt,
+                    "options": e.options,
+                    "instructions": e.instructions,
+                }))
+                sys.exit(0)
+        except Exception as e:
+            if use_json:
+                print(json.dumps({
+                    "error": type(e).__name__,
+                    "message": str(e),
+                    "traceback": traceback.format_exc(),
+                }))
+                sys.exit(1)
+            raise e
+    return wrapper

python/src/ai/chronon/click_helpers.py

@@ -17,7 +17,7 @@ def handle_compile(func):
     def wrapper(*args, **kwargs):
         if not kwargs.get("skip_compile"):
             sys.path.append(kwargs.get("repo"))
-            __compile(kwargs.get("repo"))
+            __compile(kwargs.get("repo"), force=kwargs.get("force"))
         return func(*args, **kwargs)
     return wrapper