hackagent 0.3.1__py3-none-any.whl

hackagent/router/types.py

@@ -0,0 +1,54 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+SDK-specific types for HackAgent router.
+
+These types are used internally by the SDK and are not part of the API models.
+The API uses plain strings for agent_type, but the SDK provides enums for type safety.
+"""
+
+from enum import Enum
+
+
+class AgentTypeEnum(str, Enum):
+    """
+    Enumeration of supported agent types in the HackAgent SDK.
+
+    These values correspond to the string values used in the API's agent_type field.
+
+    Endpoint Requirements by Type:
+    - GOOGLE_ADK: Google Agent Development Kit endpoint (custom protocol)
+    - LITELLM: Any LLM endpoint via LiteLLM (multi-provider support)
+    - OPENAI_SDK: OpenAI-compatible endpoint (should end with /v1 base path)
+    - LANGCHAIN: LangServe endpoint (typically /invoke or /stream)
+    - MCP: Model Context Protocol endpoint (MCP-specific protocol)
+    - A2A: Agent-to-Agent protocol endpoint (A2A-specific protocol)
+    - UNKNOWN: Unknown agent type (fallback)
+
+    Note: For OpenAI-compatible endpoints (OPENAI_SDK, LITELLM with custom endpoints),
+    provide the base URL ending in /v1 (e.g., http://localhost:8000/v1).
+    The OpenAI client will automatically append /chat/completions.
+    """
+
+    GOOGLE_ADK = "GOOGLE_ADK"
+    LITELLM = "LITELLM"
+    OPENAI_SDK = "OPENAI_SDK"
+    LANGCHAIN = "LANGCHAIN"
+    MCP = "MCP"
+    A2A = "A2A"
+    UNKNOWN = "UNKNOWN"
+
+    def __str__(self) -> str:
+        return str(self.value)

hackagent/tracking/__init__.py

@@ -0,0 +1,42 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Operation tracking and synchronization module.
+
+This module provides components for tracking pipeline operations and
+synchronizing state with the HackAgent backend API. It includes:
+
+- StepTracker: Main tracking class for managing operation lifecycle
+- track_step: Context manager for tracking individual steps
+- track_operation: Decorator for automatic operation tracking
+- TrackingContext: Shared context for tracking state
+
+The tracking system is designed to be:
+- Modular: Each component has a single responsibility
+- Reusable: Works with any attack or pipeline implementation
+- Optional: Gracefully degrades when tracking is disabled
+- Thread-safe: Safe for concurrent operations
+"""
+
+from .context import TrackingContext
+from .decorators import track_operation, track_pipeline
+from .tracker import StepTracker
+
+__all__ = [
+    "StepTracker",
+    "TrackingContext",
+    "track_operation",
+    "track_pipeline",
+]

hackagent/tracking/context.py

@@ -0,0 +1,163 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Tracking context management.
+
+This module provides the TrackingContext class for managing shared state
+across tracking operations. It acts as a lightweight container for tracking
+configuration and state that can be passed between components.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+from uuid import UUID
+
+from hackagent.client import AuthenticatedClient
+
+
+@dataclass
+class TrackingContext:
+    """
+    Shared context for operation tracking.
+
+    This class encapsulates all the state needed for tracking operations
+    and synchronizing with the backend API. It provides a clean interface
+    for passing tracking configuration between components.
+
+    Attributes:
+        client: Authenticated client for API communication
+        run_id: Server-generated run ID for this execution
+        parent_result_id: ID of the parent result record
+        logger: Logger instance for tracking operations
+        enabled: Whether tracking is enabled
+        sequence_counter: Counter for trace sequence numbers
+        metadata: Additional metadata for tracking
+
+    Example:
+        >>> context = TrackingContext(
+        ...     client=authenticated_client,
+        ...     run_id="run-123",
+        ...     parent_result_id="result-456"
+        ... )
+        >>> if context.is_enabled:
+        ...     tracker = StepTracker(context)
+    """
+
+    client: Optional[AuthenticatedClient] = None
+    run_id: Optional[str] = None
+    parent_result_id: Optional[str] = None
+    logger: Optional[logging.Logger] = None
+    sequence_counter: int = 0
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        """Initialize default logger if not provided."""
+        if self.logger is None:
+            self.logger = logging.getLogger(__name__)
+
+    @property
+    def is_enabled(self) -> bool:
+        """
+        Check if tracking is enabled.
+
+        Tracking is enabled when all required components are available:
+        client, run_id, and parent_result_id.
+
+        Returns:
+            True if tracking is enabled, False otherwise
+        """
+        return bool(
+            self.client is not None
+            and self.run_id is not None
+            and self.parent_result_id is not None
+        )
+
+    def increment_sequence(self) -> int:
+        """
+        Increment and return the sequence counter.
+
+        Returns:
+            The new sequence number
+        """
+        self.sequence_counter += 1
+        return self.sequence_counter
+
+    def get_run_uuid(self) -> Optional[UUID]:
+        """
+        Get run_id as UUID.
+
+        Returns:
+            UUID instance or None if run_id is not set
+        """
+        if self.run_id:
+            try:
+                return UUID(self.run_id)
+            except (ValueError, AttributeError):
+                self.logger.warning(f"Invalid UUID format for run_id: {self.run_id}")
+        return None
+
+    def get_result_uuid(self) -> Optional[UUID]:
+        """
+        Get parent_result_id as UUID.
+
+        Returns:
+            UUID instance or None if parent_result_id is not set
+        """
+        if self.parent_result_id:
+            try:
+                return UUID(self.parent_result_id)
+            except (ValueError, AttributeError):
+                self.logger.warning(
+                    f"Invalid UUID format for parent_result_id: {self.parent_result_id}"
+                )
+        return None
+
+    def add_metadata(self, key: str, value: Any) -> None:
+        """
+        Add metadata to the context.
+
+        Args:
+            key: Metadata key
+            value: Metadata value
+        """
+        self.metadata[key] = value
+
+    def get_metadata(self, key: str, default: Any = None) -> Any:
+        """
+        Get metadata from the context.
+
+        Args:
+            key: Metadata key
+            default: Default value if key not found
+
+        Returns:
+            Metadata value or default
+        """
+        return self.metadata.get(key, default)
+
+    @classmethod
+    def create_disabled(cls) -> "TrackingContext":
+        """
+        Create a disabled tracking context.
+
+        Returns:
+            A TrackingContext with all tracking disabled
+        """
+        return cls(
+            client=None,
+            run_id=None,
+            parent_result_id=None,
+        )

hackagent/tracking/decorators.py

@@ -0,0 +1,299 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+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.
+"""
+
+import functools
+from typing import Any, Callable, Dict, Optional, TypeVar
+
+from .tracker import StepTracker
+
+# Type variable for preserving function signatures
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+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.
+
+    Args:
+        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
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Try to get tracker from kwargs
+            tracker = kwargs.get("tracker")
+
+            # If no tracker or not a StepTracker, just run the function
+            if tracker is None or not isinstance(tracker, StepTracker):
+                return func(*args, **kwargs)
+
+            # Extract input data if extractor provided
+            input_data = None
+            if extract_input is not None:
+                try:
+                    input_data = extract_input(args, kwargs)
+                except Exception as e:
+                    tracker.logger.warning(
+                        f"Failed to extract input data for '{step_name}': {e}"
+                    )
+            else:
+                # Default extraction: look for common parameter names
+                input_data = _default_extract_input(args, kwargs)
+
+            # Extract config if extractor provided
+            config = None
+            if extract_config is not None:
+                try:
+                    config = extract_config(args, kwargs)
+                except Exception as e:
+                    tracker.logger.warning(
+                        f"Failed to extract config for '{step_name}': {e}"
+                    )
+            else:
+                # Default extraction: look for 'config' parameter
+                config = kwargs.get("config")
+
+            # Track the operation
+            with tracker.track_step(
+                step_name=step_name,
+                step_type=step_type,
+                input_data=input_data,
+                config=config,
+            ):
+                result = func(*args, **kwargs)
+                return result
+
+        return wrapper  # type: ignore
+
+    return decorator
+
+
+def _default_extract_input(args: tuple, kwargs: dict) -> Optional[Dict[str, Any]]:
+    """
+    Default input extractor for track_operation decorator.
+
+    Looks for common parameter names that might contain input data:
+    - input_df, df, data, dataframe
+    - goals, targets
+
+    Args:
+        args: Positional arguments
+        kwargs: Keyword arguments
+
+    Returns:
+        Dictionary with extracted input sample or None
+    """
+    # Try to get DataFrame-like input
+    for key in ["input_df", "df", "data", "dataframe"]:
+        if key in kwargs:
+            value = kwargs[key]
+            if hasattr(value, "head"):
+                # It's a DataFrame-like object
+                try:
+                    return {"input_sample": value.head().to_dict()}
+                except Exception:
+                    pass
+
+    # Try to get list inputs
+    for key in ["goals", "targets", "inputs"]:
+        if key in kwargs:
+            value = kwargs[key]
+            if isinstance(value, list):
+                # Sample first few items
+                sample = value[:5] if len(value) > 5 else value
+                return {key: sample}
+
+    # Try first positional argument if it's a DataFrame
+    if args and hasattr(args[0], "head"):
+        try:
+            return {"input_sample": args[0].head().to_dict()}
+        except Exception:
+            pass
+
+    return None
+
+
+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.
+
+    Args:
+        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
+    """
+
+    def decorator(cls):
+        # Store original methods
+        original_methods = {}
+
+        # Wrap all methods that have tracker parameter
+        for attr_name in dir(cls):
+            if attr_name.startswith("_"):
+                continue
+
+            attr = getattr(cls, attr_name)
+            if not callable(attr):
+                continue
+
+            # Check if method has tracker parameter
+            if hasattr(attr, "__code__"):
+                param_names = attr.__code__.co_varnames
+                if tracker_param in param_names:
+                    original_methods[attr_name] = attr
+
+        # Wrap methods to inject tracker from self
+        for method_name, original_method in original_methods.items():
+
+            @functools.wraps(original_method)
+            def wrapped_method(self, *args, _original=original_method, **kwargs):
+                # Inject tracker from self if not already provided
+                if tracker_param not in kwargs and hasattr(self, tracker_param):
+                    kwargs[tracker_param] = getattr(self, tracker_param)
+                return _original(self, *args, **kwargs)
+
+            setattr(cls, method_name, wrapped_method)
+
+        return cls
+
+    return decorator
+
+
+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.
+
+    Args:
+        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
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            # Get tracker from self if available
+            tracker = getattr(self, "tracker", None)
+
+            # If no tracker, just run the function
+            if tracker is None or not isinstance(tracker, StepTracker):
+                return func(self, *args, **kwargs)
+
+            # Extract input data
+            input_data = _default_extract_input(args, kwargs)
+
+            # Extract config (might be in kwargs or self.config)
+            config = kwargs.get("config") or getattr(self, "config", None)
+
+            # Track the operation
+            with tracker.track_step(
+                step_name=step_name,
+                step_type=step_type,
+                input_data=input_data,
+                config=config,
+            ):
+                result = func(self, *args, **kwargs)
+                return result
+
+        return wrapper  # type: ignore
+
+    return decorator