zero-agent 0.1.0__py3-none-any.whl

agentz/tools/data_tools/evaluation.py

@@ -0,0 +1,175 @@
+"""Model evaluation tool for assessing model performance."""
+
+from typing import Union, Dict, Any, Optional
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score,
+    confusion_matrix, classification_report,
+    mean_squared_error, mean_absolute_error, r2_score
+)
+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 evaluate_model(
+    ctx: RunContextWrapper[DataStore],
+    target_column: str,
+    file_path: Optional[str] = None,
+    model_type: str = "random_forest",
+    test_size: float = 0.2,
+    random_state: int = 42
+) -> Union[Dict[str, Any], str]:
+    """Evaluates machine learning model performance with comprehensive metrics.
+
+    This tool automatically uses the current dataset from the pipeline context.
+    A file_path can optionally be provided to evaluate on a different dataset.
+
+    Args:
+        ctx: Pipeline context wrapper for accessing the data store
+        target_column: Name of the target column to predict
+        file_path: Optional path to dataset file. If not provided, uses current dataset.
+        model_type: Type of model to evaluate (random_forest, decision_tree, etc.)
+        test_size: Proportion of data to use for testing (default: 0.2)
+        random_state: Random seed for reproducibility (default: 42)
+
+    Returns:
+        Dictionary containing:
+            - problem_type: "classification" or "regression"
+            - metrics: Performance metrics
+            - confusion_matrix: Confusion matrix (for classification)
+            - classification_report: Detailed classification report
+            - cross_validation: Cross-validation results
+            - error_analysis: Error distribution analysis
+        Or error message string if evaluation 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("Evaluating model on 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=True, data_store=data_store)
+            logger.info(f"Evaluating model on dataset from: {file_path}")
+
+        if target_column not in df.columns:
+            return f"Target column '{target_column}' not found in dataset"
+
+        # Separate features and target
+        X = df.drop(columns=[target_column])
+        y = df[target_column]
+
+        # Handle categorical features
+        X = pd.get_dummies(X, drop_first=True)
+
+        # Determine problem type
+        is_classification = y.dtype == 'object' or y.nunique() < 20
+
+        # Encode target if categorical
+        original_labels = None
+        if is_classification and y.dtype == 'object':
+            from sklearn.preprocessing import LabelEncoder
+            le = LabelEncoder()
+            original_labels = le.classes_
+            y = le.fit_transform(y)
+
+        # Split data
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=test_size, random_state=random_state
+        )
+
+        # Train model
+        if is_classification:
+            model = RandomForestClassifier(random_state=random_state, n_estimators=100)
+        else:
+            model = RandomForestRegressor(random_state=random_state, n_estimators=100)
+
+        model.fit(X_train, y_train)
+        y_pred = model.predict(X_test)
+
+        result = {
+            "problem_type": "classification" if is_classification else "regression",
+        }
+
+        if is_classification:
+            # Classification metrics
+            metrics = {
+                "accuracy": float(accuracy_score(y_test, y_pred)),
+                "precision": float(precision_score(y_test, y_pred, average='weighted', zero_division=0)),
+                "recall": float(recall_score(y_test, y_pred, average='weighted', zero_division=0)),
+                "f1_score": float(f1_score(y_test, y_pred, average='weighted', zero_division=0)),
+            }
+            result["metrics"] = metrics
+
+            # Confusion matrix
+            cm = confusion_matrix(y_test, y_pred)
+            result["confusion_matrix"] = cm.tolist()
+
+            # Classification report
+            if original_labels is not None:
+                target_names = [str(label) for label in original_labels]
+            else:
+                target_names = [str(i) for i in sorted(np.unique(y))]
+
+            class_report = classification_report(y_test, y_pred, target_names=target_names, output_dict=True, zero_division=0)
+            result["classification_report"] = class_report
+
+            # Per-class accuracy
+            per_class_accuracy = {}
+            for i, label in enumerate(target_names):
+                mask = y_test == i
+                if mask.sum() > 0:
+                    per_class_accuracy[label] = float(accuracy_score(y_test[mask], y_pred[mask]))
+            result["per_class_accuracy"] = per_class_accuracy
+
+        else:
+            # Regression metrics
+            metrics = {
+                "r2_score": float(r2_score(y_test, y_pred)),
+                "mean_squared_error": float(mean_squared_error(y_test, y_pred)),
+                "root_mean_squared_error": float(np.sqrt(mean_squared_error(y_test, y_pred))),
+                "mean_absolute_error": float(mean_absolute_error(y_test, y_pred)),
+                "mean_absolute_percentage_error": float(np.mean(np.abs((y_test - y_pred) / y_test)) * 100),
+            }
+            result["metrics"] = metrics
+
+            # Error analysis
+            errors = y_test - y_pred
+            error_analysis = {
+                "mean_error": float(np.mean(errors)),
+                "std_error": float(np.std(errors)),
+                "min_error": float(np.min(errors)),
+                "max_error": float(np.max(errors)),
+                "median_error": float(np.median(errors)),
+            }
+            result["error_analysis"] = error_analysis
+
+        # Cross-validation
+        cv_scores = cross_val_score(model, X, y, cv=5)
+        result["cross_validation"] = {
+            "scores": cv_scores.tolist(),
+            "mean": float(cv_scores.mean()),
+            "std": float(cv_scores.std()),
+        }
+
+        # Feature importance
+        if hasattr(model, 'feature_importances_'):
+            importance_dict = dict(zip(X.columns, model.feature_importances_))
+            sorted_importance = dict(sorted(importance_dict.items(), key=lambda x: x[1], reverse=True))
+            result["feature_importance"] = {k: float(v) for k, v in list(sorted_importance.items())[:10]}
+
+        return result
+
+    except Exception as e:
+        return f"Error evaluating model: {str(e)}"

agentz/tools/data_tools/helpers.py

@@ -0,0 +1,120 @@
+"""Helper utilities for data tools to access cached DataFrames."""
+
+from typing import Optional
+from pathlib import Path
+import pandas as pd
+from agentz.context.data_store import DataStore
+from loguru import logger
+
+
+def get_dataframe(file_path: str, prefer_preprocessed: bool = False, data_store: Optional[DataStore] = None) -> Optional[pd.DataFrame]:
+    """Get a DataFrame from cache or load from file.
+
+    Args:
+        file_path: Path to the dataset file
+        prefer_preprocessed: If True, try to get preprocessed version first
+        data_store: Pipeline data store instance
+
+    Returns:
+        DataFrame or None if not found
+    """
+    file_path = Path(file_path)
+
+    if not data_store:
+        return None
+
+    # Try preprocessed first if requested
+    if prefer_preprocessed:
+        preprocessed_key = f"preprocessed:{file_path.resolve()}"
+        if data_store.has(preprocessed_key):
+            logger.info(f"Using cached preprocessed DataFrame for {file_path}")
+            return data_store.get(preprocessed_key)
+
+    # Try regular DataFrame
+    cache_key = f"dataframe:{file_path.resolve()}"
+    if data_store.has(cache_key):
+        logger.info(f"Using cached DataFrame for {file_path}")
+        return data_store.get(cache_key)
+
+    return None
+
+
+def load_or_get_dataframe(file_path: str, prefer_preprocessed: bool = False, data_store: Optional[DataStore] = None) -> pd.DataFrame:
+    """Get DataFrame from cache or load from file, with fallback loading.
+
+    Args:
+        file_path: Path to the dataset file
+        prefer_preprocessed: If True, try to get preprocessed version first
+        data_store: Pipeline data store instance
+
+    Returns:
+        DataFrame
+
+    Raises:
+        FileNotFoundError: If file doesn't exist and not in cache
+        ValueError: If file format is not supported
+    """
+    # Try cache first
+    df = get_dataframe(file_path, prefer_preprocessed, data_store)
+    if df is not None:
+        return df
+
+    # Load from file
+    file_path = Path(file_path)
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    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:
+        raise ValueError(f"Unsupported file format: {file_path.suffix}")
+
+    # Cache it for future use
+    if data_store:
+        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}")
+
+    return df
+
+
+def cache_object(key: str, obj: any, data_type: str = None, metadata: dict = None, data_store: Optional[DataStore] = None) -> None:
+    """Cache an object in the pipeline data store.
+
+    Args:
+        key: Cache key
+        obj: Object to cache
+        data_type: Type descriptor (e.g., 'model', 'scaler')
+        metadata: Optional metadata
+        data_store: Pipeline data store instance
+    """
+    if data_store:
+        data_store.set(key, obj, data_type=data_type, metadata=metadata)
+        logger.info(f"Cached {data_type or 'object'} with key: {key}")
+
+
+def get_cached_object(key: str, data_store: Optional[DataStore] = None) -> Optional[any]:
+    """Get a cached object from the pipeline data store.
+
+    Args:
+        key: Cache key
+        data_store: Pipeline data store instance
+
+    Returns:
+        Cached object or None
+    """
+    if data_store and data_store.has(key):
+        logger.info(f"Retrieved cached object with key: {key}")
+        return data_store.get(key)
+    return None

agentz/tools/data_tools/model_training.py

@@ -0,0 +1,192 @@
+"""Model training tool for training machine learning models."""
+
+from typing import Union, Dict, Any, Optional, List
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
+from sklearn.metrics import accuracy_score, mean_squared_error, r2_score
+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, cache_object
+from loguru import logger
+
+
+@function_tool
+async def train_model(
+    ctx: RunContextWrapper[DataStore],
+    target_column: str,
+    file_path: Optional[str] = None,
+    model_type: str = "auto",
+    test_size: float = 0.2,
+    random_state: int = 42
+) -> Union[Dict[str, Any], str]:
+    """Trains machine learning models on a dataset.
+
+    This tool automatically uses the current dataset from the pipeline context.
+    A file_path can optionally be provided to train on a different dataset.
+
+    Args:
+        ctx: Pipeline context wrapper for accessing the data store
+        target_column: Name of the target column to predict
+        file_path: Optional path to dataset file. If not provided, uses current dataset.
+        model_type: Type of model to train. Options:
+            - "auto": Automatically detect and use best model
+            - "random_forest": Random Forest
+            - "logistic_regression": Logistic Regression (classification)
+            - "linear_regression": Linear Regression (regression)
+            - "decision_tree": Decision Tree
+        test_size: Proportion of data to use for testing (default: 0.2)
+        random_state: Random seed for reproducibility (default: 42)
+
+    Returns:
+        Dictionary containing:
+            - model_type: Type of model trained
+            - problem_type: "classification" or "regression"
+            - train_score: Training score
+            - test_score: Testing score
+            - cross_val_scores: Cross-validation scores (mean and std)
+            - feature_importance: Feature importance scores (if available)
+            - predictions_sample: Sample of predictions vs actual values
+        Or error message string if training 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("Training model on 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=True, data_store=data_store)
+            logger.info(f"Training model on dataset from: {file_path}")
+
+        if target_column not in df.columns:
+            return f"Target column '{target_column}' not found in dataset"
+
+        # Separate features and target
+        X = df.drop(columns=[target_column])
+        y = df[target_column]
+
+        # Handle categorical features in X
+        X = pd.get_dummies(X, drop_first=True)
+
+        # Determine problem type
+        is_classification = y.dtype == 'object' or y.nunique() < 20
+
+        # Encode target if categorical
+        if is_classification and y.dtype == 'object':
+            from sklearn.preprocessing import LabelEncoder
+            le = LabelEncoder()
+            y = le.fit_transform(y)
+
+        # Split data
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=test_size, random_state=random_state
+        )
+
+        # Select model
+        if model_type == "auto":
+            if is_classification:
+                model = RandomForestClassifier(random_state=random_state, n_estimators=100)
+                model_name = "Random Forest Classifier"
+            else:
+                model = RandomForestRegressor(random_state=random_state, n_estimators=100)
+                model_name = "Random Forest Regressor"
+        elif model_type == "random_forest":
+            if is_classification:
+                model = RandomForestClassifier(random_state=random_state, n_estimators=100)
+                model_name = "Random Forest Classifier"
+            else:
+                model = RandomForestRegressor(random_state=random_state, n_estimators=100)
+                model_name = "Random Forest Regressor"
+        elif model_type == "logistic_regression":
+            model = LogisticRegression(random_state=random_state, max_iter=1000)
+            model_name = "Logistic Regression"
+        elif model_type == "linear_regression":
+            model = LinearRegression()
+            model_name = "Linear Regression"
+        elif model_type == "decision_tree":
+            if is_classification:
+                model = DecisionTreeClassifier(random_state=random_state)
+                model_name = "Decision Tree Classifier"
+            else:
+                model = DecisionTreeRegressor(random_state=random_state)
+                model_name = "Decision Tree Regressor"
+        else:
+            return f"Unknown model type: {model_type}"
+
+        # Train model
+        model.fit(X_train, y_train)
+
+        # Evaluate
+        train_pred = model.predict(X_train)
+        test_pred = model.predict(X_test)
+
+        if is_classification:
+            train_score = accuracy_score(y_train, train_pred)
+            test_score = accuracy_score(y_test, test_pred)
+            metric_name = "accuracy"
+        else:
+            train_score = r2_score(y_train, train_pred)
+            test_score = r2_score(y_test, test_pred)
+            metric_name = "r2_score"
+
+        # Cross-validation
+        cv_scores = cross_val_score(model, X, y, cv=5)
+
+        # Feature importance
+        feature_importance = {}
+        if hasattr(model, 'feature_importances_'):
+            importance_dict = dict(zip(X.columns, model.feature_importances_))
+            # Sort by importance
+            feature_importance = dict(sorted(importance_dict.items(), key=lambda x: x[1], reverse=True))
+
+        # Sample predictions
+        predictions_sample = []
+        for i in range(min(10, len(y_test))):
+            predictions_sample.append({
+                "actual": float(y_test.iloc[i]) if hasattr(y_test, 'iloc') else float(y_test[i]),
+                "predicted": float(test_pred[i]),
+            })
+
+        # Cache the trained model for reuse
+        file_path_obj = Path(file_path) if file_path else Path("model")
+        model_key = f"model:{file_path_obj.resolve()}"
+        cache_object(
+            model_key,
+            model,
+            data_type="model",
+            data_store=data_store,
+            metadata={
+                "file_path": str(file_path),
+                "model_type": model_name,
+                "target_column": target_column,
+                "test_score": float(test_score)
+            }
+        )
+
+        result = {
+            "model_type": model_name,
+            "problem_type": "classification" if is_classification else "regression",
+            "train_score": float(train_score),
+            "test_score": float(test_score),
+            "metric": metric_name,
+            "cross_val_mean": float(cv_scores.mean()),
+            "cross_val_std": float(cv_scores.std()),
+            "feature_importance": {k: float(v) for k, v in list(feature_importance.items())[:10]},  # Top 10
+            "predictions_sample": predictions_sample,
+            "train_size": len(X_train),
+            "test_size": len(X_test),
+        }
+
+        return result
+
+    except Exception as e:
+        return f"Error training model: {str(e)}"