contextagent 0.1.0__py3-none-any.whl

agentz/runner/workflow.py

@@ -0,0 +1,250 @@
+"""Workflow control flow helpers for pipeline execution."""
+
+import asyncio
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
+
+
+class WorkflowHelpers:
+    """Collection of workflow control flow helpers.
+
+    These helpers provide common patterns for pipeline execution:
+    - Iterative loops
+    - Custom groups
+    - Parallel execution
+    - Conditional execution
+    - Loop-until patterns
+    """
+
+    def __init__(
+        self,
+        iteration_manager: Any,
+        hook_registry: Any,
+        start_group_callback: Optional[Callable] = None,
+        end_group_callback: Optional[Callable] = None,
+        update_printer_callback: Optional[Callable] = None,
+    ):
+        """Initialize workflow helpers.
+
+        Args:
+            iteration_manager: IterationManager instance
+            hook_registry: HookRegistry instance
+            start_group_callback: Optional callback(group_id, title, border_style)
+            end_group_callback: Optional callback(group_id, is_done)
+            update_printer_callback: Optional callback(key, message, is_done, group_id)
+        """
+        self.iteration_manager = iteration_manager
+        self.hook_registry = hook_registry
+        self.start_group = start_group_callback
+        self.end_group = end_group_callback
+        self.update_printer = update_printer_callback
+
+    async def run_iterative_loop(
+        self,
+        iteration_body: Callable[[Any, str], Awaitable[Any]],
+        final_body: Optional[Callable[[str], Awaitable[Any]]] = None,
+        should_continue: Optional[Callable[[], bool]] = None,
+    ) -> Any:
+        """Execute standard iterative loop pattern.
+
+        Args:
+            iteration_body: Async function(iteration, group_id) -> result
+            final_body: Optional async function(final_group_id) -> result
+            should_continue: Optional custom condition (default: iteration_manager.should_continue_iteration)
+
+        Returns:
+            Result from final_body if provided, else None
+
+        Example:
+            async def my_iteration(iteration, group):
+                observations = await observe_agent(...)
+                evaluations = await evaluate_agent(...)
+                await route_and_execute(evaluations, group)
+
+            async def my_final(group):
+                return await writer_agent(...)
+
+            result = await helpers.run_iterative_loop(
+                iteration_body=my_iteration,
+                final_body=my_final
+            )
+        """
+        should_continue_fn = should_continue or self.iteration_manager.should_continue_iteration
+
+        while should_continue_fn():
+            iteration, group_id = self.iteration_manager.begin_iteration()
+
+            await self.hook_registry.trigger(
+                "before_iteration",
+                context=self.iteration_manager.context,
+                iteration=iteration,
+                group_id=group_id
+            )
+
+            try:
+                await iteration_body(iteration, group_id)
+            finally:
+                await self.hook_registry.trigger(
+                    "after_iteration",
+                    context=self.iteration_manager.context,
+                    iteration=iteration,
+                    group_id=group_id
+                )
+                self.iteration_manager.end_iteration(group_id)
+
+            # Check if state indicates completion
+            context = self.iteration_manager.context
+            if hasattr(context, 'state') and context.state and context.state.complete:
+                break
+
+        result = None
+        if final_body:
+            final_group = self.iteration_manager.start_final_group()
+            result = await final_body(final_group)
+            self.iteration_manager.end_final_group(final_group)
+
+        return result
+
+    async def run_custom_group(
+        self,
+        group_id: str,
+        title: str,
+        body: Callable[[], Awaitable[Any]],
+        border_style: str = "white",
+    ) -> Any:
+        """Execute code within a custom printer group.
+
+        Args:
+            group_id: Unique group identifier
+            title: Display title for the group
+            body: Async function to execute within group
+            border_style: Border color for printer
+
+        Returns:
+            Result from body()
+
+        Example:
+            exploration = await helpers.run_custom_group(
+                "exploration",
+                "Exploration Phase",
+                self._explore
+            )
+
+            analysis = await helpers.run_custom_group(
+                "analysis",
+                "Deep Analysis",
+                lambda: self._analyze(exploration)
+            )
+        """
+        if self.start_group:
+            self.start_group(group_id, title=title, border_style=border_style)
+        try:
+            result = await body()
+            return result
+        finally:
+            if self.end_group:
+                self.end_group(group_id, is_done=True)
+
+    async def run_parallel_steps(
+        self,
+        steps: Dict[str, Callable[[], Awaitable[Any]]],
+        group_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Execute multiple steps in parallel.
+
+        Args:
+            steps: Dict mapping step_name -> async callable
+            group_id: Optional group to nest steps in
+
+        Returns:
+            Dict mapping step_name -> result
+
+        Example:
+            results = await helpers.run_parallel_steps({
+                "data_loading": self.load_data,
+                "validation": self.validate_inputs,
+                "model_init": self.initialize_models,
+            })
+
+            data = results["data_loading"]
+        """
+        async def run_step(name: str, fn: Callable):
+            key = f"{group_id}:{name}" if group_id else name
+            if self.update_printer:
+                self.update_printer(key, f"Running {name}...", group_id=group_id)
+            result = await fn()
+            if self.update_printer:
+                self.update_printer(key, f"Completed {name}", is_done=True, group_id=group_id)
+            return name, result
+
+        tasks = [run_step(name, fn) for name, fn in steps.items()]
+        completed = await asyncio.gather(*tasks)
+        return dict(completed)
+
+    async def run_if(
+        self,
+        condition: Union[bool, Callable[[], bool]],
+        body: Callable[[], Awaitable[Any]],
+        else_body: Optional[Callable[[], Awaitable[Any]]] = None,
+    ) -> Any:
+        """Conditional execution helper.
+
+        Args:
+            condition: Boolean or callable returning bool
+            body: Execute if condition is True
+            else_body: Optional execute if condition is False
+
+        Returns:
+            Result from executed body
+
+        Example:
+            initial = await quick_check()
+
+            return await helpers.run_if(
+                condition=initial.needs_deep_analysis,
+                body=lambda: deep_analysis(initial),
+                else_body=lambda: simple_report(initial)
+            )
+        """
+        cond_result = condition() if callable(condition) else condition
+        if cond_result:
+            return await body()
+        elif else_body:
+            return await else_body()
+        return None
+
+    async def run_until(
+        self,
+        condition: Callable[[], bool],
+        body: Callable[[int], Awaitable[Any]],
+        max_iterations: Optional[int] = None,
+    ) -> List[Any]:
+        """Execute body repeatedly until condition is met.
+
+        Args:
+            condition: Callable returning True to stop
+            body: Async function(iteration_number) -> result
+            max_iterations: Optional max iterations (default: unlimited)
+
+        Returns:
+            List of results from each iteration
+
+        Example:
+            results = await helpers.run_until(
+                condition=lambda: context.state.complete,
+                body=self._exploration_step,
+                max_iterations=10
+            )
+            return aggregate(results)
+        """
+        results = []
+        iteration = 0
+
+        while not condition():
+            if max_iterations and iteration >= max_iterations:
+                break
+
+            result = await body(iteration)
+            results.append(result)
+            iteration += 1
+
+        return results

agentz/tools/__init__.py

@@ -0,0 +1,20 @@
+"""Tools for agent workflows."""
+
+# Re-export data tools for backward compatibility
+from .data_tools import (
+    load_dataset,
+    analyze_data,
+    preprocess_data,
+    train_model,
+    evaluate_model,
+    create_visualization,
+)
+
+__all__ = [
+    "load_dataset",
+    "analyze_data",
+    "preprocess_data",
+    "train_model",
+    "evaluate_model",
+    "create_visualization",
+]

agentz/tools/data_tools/__init__.py

@@ -0,0 +1,17 @@
+"""Data science tools for data analysis, preprocessing, modeling, and visualization."""
+
+from .data_loading import load_dataset
+from .data_analysis import analyze_data
+from .preprocessing import preprocess_data
+from .model_training import train_model
+from .evaluation import evaluate_model
+from .visualization import create_visualization
+
+__all__ = [
+    "load_dataset",
+    "analyze_data",
+    "preprocess_data",
+    "train_model",
+    "evaluate_model",
+    "create_visualization",
+]

agentz/tools/data_tools/data_analysis.py

@@ -0,0 +1,152 @@
+"""Data analysis tool for exploratory data analysis and statistical analysis."""
+
+from typing import Union, Dict, Any, Optional
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from agents import function_tool
+from agents.run_context import RunContextWrapper
+from agentz.context.data_store import DataStore
+from .helpers import load_or_get_dataframe
+from loguru import logger
+
+
+@function_tool
+async def analyze_data(ctx: RunContextWrapper[DataStore], file_path: Optional[str] = None, target_column: str = None) -> Union[Dict[str, Any], str]:
+    """Performs comprehensive exploratory data analysis on a dataset.
+
+    This tool automatically uses the current dataset from the pipeline context.
+    A file_path can optionally be provided to analyze a different dataset.
+
+    Args:
+        ctx: Pipeline context wrapper for accessing the data store
+        file_path: Optional path to dataset file. If not provided, uses current dataset.
+        target_column: Optional target column for correlation analysis
+
+    Returns:
+        Dictionary containing:
+            - distributions: Distribution statistics for each column
+            - correlations: Correlation matrix for numerical columns
+            - outliers: Outlier detection results using IQR method
+            - patterns: Identified patterns and insights
+            - recommendations: Data quality and preprocessing recommendations
+        Or error message string if analysis fails
+    """
+    try:
+        # Get DataFrame - either from file_path or current dataset
+        data_store = ctx.context
+        if file_path is None:
+            if data_store and data_store.has("current_dataset"):
+                df = data_store.get("current_dataset")
+                logger.info("Analyzing current dataset from pipeline context")
+            else:
+                return "Error: No dataset loaded. Please load a dataset first using the load_dataset tool."
+        else:
+            df = load_or_get_dataframe(file_path, prefer_preprocessed=False, data_store=data_store)
+            logger.info(f"Analyzing dataset from: {file_path}")
+
+        result = {}
+
+        # Distribution analysis
+        distributions = {}
+        for col in df.columns:
+            if pd.api.types.is_numeric_dtype(df[col]):
+                distributions[col] = {
+                    "mean": float(df[col].mean()),
+                    "median": float(df[col].median()),
+                    "std": float(df[col].std()),
+                    "min": float(df[col].min()),
+                    "max": float(df[col].max()),
+                    "q25": float(df[col].quantile(0.25)),
+                    "q75": float(df[col].quantile(0.75)),
+                    "skewness": float(df[col].skew()),
+                    "kurtosis": float(df[col].kurtosis()),
+                }
+            else:
+                distributions[col] = {
+                    "unique_values": int(df[col].nunique()),
+                    "top_value": str(df[col].mode()[0]) if not df[col].mode().empty else None,
+                    "top_frequency": int(df[col].value_counts().iloc[0]) if len(df[col].value_counts()) > 0 else 0,
+                }
+        result["distributions"] = distributions
+
+        # Correlation analysis
+        numeric_df = df.select_dtypes(include=['number'])
+        if not numeric_df.empty:
+            corr_matrix = numeric_df.corr()
+            result["correlations"] = corr_matrix.to_dict()
+
+            if target_column and target_column in corr_matrix.columns:
+                target_corr = corr_matrix[target_column].drop(target_column).sort_values(ascending=False)
+                result["target_correlations"] = target_corr.to_dict()
+
+        # Outlier detection using IQR method
+        outliers = {}
+        for col in numeric_df.columns:
+            Q1 = df[col].quantile(0.25)
+            Q3 = df[col].quantile(0.75)
+            IQR = Q3 - Q1
+            lower_bound = Q1 - 1.5 * IQR
+            upper_bound = Q3 + 1.5 * IQR
+            outlier_count = ((df[col] < lower_bound) | (df[col] > upper_bound)).sum()
+            outliers[col] = {
+                "count": int(outlier_count),
+                "percentage": float(outlier_count / len(df) * 100),
+                "lower_bound": float(lower_bound),
+                "upper_bound": float(upper_bound),
+            }
+        result["outliers"] = outliers
+
+        # Pattern identification
+        patterns = []
+
+        # High correlation patterns
+        if "correlations" in result:
+            for col1 in corr_matrix.columns:
+                for col2 in corr_matrix.columns:
+                    if col1 < col2:  # Avoid duplicates
+                        corr_val = corr_matrix.loc[col1, col2]
+                        if abs(corr_val) > 0.7:
+                            patterns.append(f"Strong correlation ({corr_val:.2f}) between {col1} and {col2}")
+
+        # Missing data patterns
+        missing_cols = [col for col in df.columns if df[col].isnull().sum() > 0]
+        if missing_cols:
+            patterns.append(f"Missing data detected in {len(missing_cols)} columns: {', '.join(missing_cols[:5])}")
+
+        # Outlier patterns
+        high_outlier_cols = [col for col, info in outliers.items() if info['percentage'] > 5]
+        if high_outlier_cols:
+            patterns.append(f"High outlier percentage (>5%) in columns: {', '.join(high_outlier_cols)}")
+
+        result["patterns"] = patterns
+
+        # Recommendations
+        recommendations = []
+
+        if missing_cols:
+            recommendations.append("Consider imputation strategies for missing values")
+
+        if high_outlier_cols:
+            recommendations.append("Review and handle outliers before modeling")
+
+        # Check for imbalanced categorical columns
+        for col in df.select_dtypes(include=['object']).columns:
+            value_counts = df[col].value_counts()
+            if len(value_counts) > 1:
+                imbalance_ratio = value_counts.iloc[0] / value_counts.iloc[-1]
+                if imbalance_ratio > 10:
+                    recommendations.append(f"Column '{col}' shows class imbalance (ratio: {imbalance_ratio:.1f})")
+
+        # Check for constant or near-constant columns
+        for col in df.columns:
+            unique_ratio = df[col].nunique() / len(df)
+            if unique_ratio < 0.01 and df[col].nunique() > 1:
+                recommendations.append(f"Column '{col}' has very low variance, consider removing")
+
+        result["recommendations"] = recommendations
+
+        return result
+
+    except Exception as e:
+        return f"Error analyzing dataset: {str(e)}"

agentz/tools/data_tools/data_loading.py

@@ -0,0 +1,92 @@
+"""Data loading tool for loading and inspecting datasets."""
+
+from typing import Union, Dict, Any
+from pathlib import Path
+import pandas as pd
+from agents import function_tool
+from agents.run_context import RunContextWrapper
+from agentz.context.data_store import DataStore
+from loguru import logger
+
+
+@function_tool
+async def load_dataset(ctx: RunContextWrapper[DataStore], file_path: str) -> Union[Dict[str, Any], str]:
+    """Loads a dataset and provides comprehensive inspection information.
+
+    This tool caches the loaded DataFrame in the pipeline data store so other
+    tools can reuse it without reloading from disk.
+
+    Args:
+        ctx: Pipeline context wrapper for accessing the data store
+        file_path: Path to the dataset file (CSV, JSON, Excel, etc.)
+
+    Returns:
+        Dictionary containing:
+            - shape: Tuple of (rows, columns)
+            - columns: List of column names
+            - dtypes: Dictionary of column data types
+            - missing_values: Dictionary of missing value counts per column
+            - sample_data: First 5 rows as dictionary
+            - summary_stats: Statistical summary for numerical columns
+            - memory_usage: Memory usage information
+        Or error message string if loading fails
+    """
+    try:
+        file_path = Path(file_path)
+
+        if not file_path.exists():
+            return f"File not found: {file_path}"
+
+        # Load based on file extension
+        if file_path.suffix.lower() == '.csv':
+            df = pd.read_csv(file_path)
+        elif file_path.suffix.lower() in ['.xlsx', '.xls']:
+            df = pd.read_excel(file_path)
+        elif file_path.suffix.lower() == '.json':
+            df = pd.read_json(file_path)
+        elif file_path.suffix.lower() == '.parquet':
+            df = pd.read_parquet(file_path)
+        else:
+            return f"Unsupported file format: {file_path.suffix}"
+
+        # Store DataFrame in data store for reuse by other tools
+        data_store = ctx.context
+        if data_store is not None:
+            # Store with file path key for backward compatibility
+            cache_key = f"dataframe:{file_path.resolve()}"
+            data_store.set(
+                cache_key,
+                df,
+                data_type="dataframe",
+                metadata={"file_path": str(file_path), "shape": df.shape}
+            )
+            logger.info(f"Cached DataFrame from {file_path} with key: {cache_key}")
+
+            # Also set as the current active dataset
+            data_store.set(
+                "current_dataset",
+                df,
+                data_type="dataframe",
+                metadata={"file_path": str(file_path), "shape": df.shape, "source": "loaded"}
+            )
+            logger.info(f"Set as current dataset for pipeline")
+
+        # Gather comprehensive information
+        result = {
+            "file_path": str(file_path),
+            "shape": df.shape,
+            "columns": df.columns.tolist(),
+            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
+            "missing_values": df.isnull().sum().to_dict(),
+            "missing_percentage": (df.isnull().sum() / len(df) * 100).to_dict(),
+            "sample_data": df.head(5).to_dict(orient='records'),
+            "summary_stats": df.describe().to_dict() if not df.select_dtypes(include=['number']).empty else {},
+            "memory_usage": df.memory_usage(deep=True).to_dict(),
+            "total_memory_mb": df.memory_usage(deep=True).sum() / 1024 / 1024,
+            "duplicate_rows": int(df.duplicated().sum()),
+        }
+
+        return result
+
+    except Exception as e:
+        return f"Error loading dataset: {str(e)}"