Skip to main content

hackagent.router.tracking.decorators

Decorators for automatic operation tracking.

This module provides decorator functions that can be applied to functions or methods to automatically track their execution. Decorators offer a declarative way to add tracking without modifying function bodies.

track_operation

def track_operation(
step_name: str,
step_type: str,
extract_input: Optional[Callable[[Any, Any], Dict[str, Any]]] = None,
extract_config: Optional[Callable[[Any, Any], Dict[str, Any]]] = None
) -> Callable[[F], F]

Decorator for automatic operation tracking.

This decorator wraps a function to automatically track its execution using a StepTracker. It looks for a 'tracker' parameter in the function arguments and uses it if available.

The decorator is flexible and can extract input data and configuration using custom extractor functions, allowing it to work with any function signature.

Arguments:

  • step_name - Human-readable name for the operation
  • step_type - Step type identifier (e.g., "STEP1_GENERATE")
  • extract_input - Optional function to extract input data from args/kwargs
  • extract_config - Optional function to extract config from args/kwargs

Returns:

Decorated function with automatic tracking

Example:

>>> @track_operation("Generate Prefixes", "STEP1_GENERATE") ... def generate_prefixes(goals, config, tracker=None): ... # Function logic ... return results

>>> # With custom extractors >>> def get_input(args, kwargs): ... return {"goals": kwargs.get("goals", [])} >>> >>> @track_operation( ... "Process Data", ... "STEP2_PROCESS", ... extract_input=get_input ... ) ... def process_data(data, config, tracker=None): ... return processed_data

track_pipeline

def track_pipeline(tracker_param: str = "tracker")

Class decorator for automatic pipeline tracking.

This decorator can be applied to a class to make all its methods automatically aware of a tracker instance. It's useful for pipeline classes where multiple methods should be tracked.

Arguments:

  • tracker_param - Name of the parameter that contains the tracker

Returns:

Decorated class with tracking support

Example:

>>> @track_pipeline(tracker_param="tracker") ... class MyPipeline: ... def init(self, tracker=None): ... self.tracker = tracker ... ... @track_operation("Step 1", "STEP1") ... def step1(self, data, tracker=None): ... return processed_data ... ... @track_operation("Step 2", "STEP2") ... def step2(self, data, tracker=None): ... return final_data

>>> # All methods will automatically use self.tracker >>> pipeline = MyPipeline(tracker=my_tracker) >>> pipeline.step1(data) # Automatically tracked

track_method

def track_method(step_name: str, step_type: str)

Method decorator that automatically uses self.tracker.

This is a specialized version of track_operation designed for class methods. It automatically looks for self.tracker and uses it for tracking.

Arguments:

  • step_name - Human-readable name for the operation
  • step_type - Step type identifier

Returns:

Decorated method with automatic tracking

Example:

>>> class Pipeline: ... def init(self, tracker): ... self.tracker = tracker ... ... @track_method("Generate Data", "STEP1") ... def generate(self, goals): ... return generated_data ... ... @track_method("Process Data", "STEP2") ... def process(self, data): ... return processed_data