ds-agent-cli 0.1.0

package/src/tools/agent_tool_mapping.py

@@ -0,0 +1,326 @@
+"""
+Agent-Specific Tool Mapping
+Maps specialist agents to their relevant tools for dynamic loading.
+"""
+
+from ds_agent.registry.tools_registry import TOOLS
+
+# Define tool categories and their tools
+TOOL_CATEGORIES = {
+    "profiling": [
+        "profile_dataset",
+        "detect_data_quality_issues",
+        "analyze_correlations",
+        "get_smart_summary",
+    ],
+    "cleaning": [
+        "clean_missing_values",
+        "handle_outliers",
+        "fix_data_types",
+        "force_numeric_conversion",
+        "smart_type_inference",
+        "remove_duplicates",
+    ],
+    "feature_engineering": [
+        "create_time_features",
+        "encode_categorical",
+        "create_interaction_features",
+        "create_ratio_features",
+        "create_statistical_features",
+        "create_log_features",
+        "create_binned_features",
+        "create_aggregation_features",
+        "auto_feature_engineering",
+    ],
+    "visualization": [
+        "generate_interactive_scatter",
+        "generate_interactive_histogram",
+        "generate_interactive_box_plots",
+        "generate_interactive_correlation_heatmap",
+        "generate_interactive_time_series",
+        "generate_plotly_dashboard",
+        "generate_eda_plots",
+        "generate_combined_eda_report",
+    ],
+    "modeling": [
+        "train_baseline_models",
+        "train_with_autogluon",
+        "predict_with_autogluon",
+        "optimize_autogluon_model",
+        "analyze_autogluon_model",
+        "extend_autogluon_training",
+        "train_multilabel_autogluon",
+        "hyperparameter_tuning",
+        "perform_cross_validation",
+        "train_ensemble_models",
+        "auto_ml_pipeline",
+        "evaluate_model_performance",
+    ],
+    "time_series": [
+        "detect_seasonality",
+        "decompose_time_series",
+        "forecast_arima",
+        "forecast_prophet",
+        "forecast_with_autogluon",
+        "backtest_timeseries",
+        "analyze_timeseries_model",
+        "detect_anomalies_time_series",
+    ],
+    "nlp": [
+        "extract_entities",
+        "sentiment_analysis",
+        "topic_modeling",
+        "text_classification",
+        "text_preprocessing",
+    ],
+    "computer_vision": [
+        "image_classification",
+        "object_detection",
+        "image_preprocessing",
+    ],
+    "business_intelligence": [
+        "calculate_kpis",
+        "trend_analysis",
+        "cohort_analysis",
+        "churn_prediction",
+    ],
+    "production": [
+        "export_model_to_onnx",
+        "generate_inference_code",
+        "create_model_documentation",
+        "validate_model_drift",
+    ],
+    "code_execution": [
+        "execute_python_code",
+        "debug_code",
+    ]
+}
+
+# Map specialist agents to their relevant tool categories
+AGENT_TOOL_MAPPING = {
+    "data_quality_agent": {
+        "categories": ["profiling", "cleaning"],
+        "description": "Focuses on data profiling, quality assessment, and cleaning operations"
+    },
+    "preprocessing_agent": {
+        "categories": ["cleaning", "feature_engineering", "profiling"],
+        "description": "Handles data cleaning, transformation, and feature engineering"
+    },
+    "visualization_agent": {
+        "categories": ["visualization", "profiling"],
+        "description": "Creates charts, plots, and interactive dashboards"
+    },
+    "modeling_agent": {
+        "categories": ["modeling", "feature_engineering", "profiling"],
+        "description": "Trains, tunes, and evaluates machine learning models"
+    },
+    "time_series_agent": {
+        "categories": ["time_series", "profiling", "visualization"],
+        "description": "Specializes in time series analysis and forecasting"
+    },
+    "nlp_agent": {
+        "categories": ["nlp", "profiling", "visualization"],
+        "description": "Natural language processing and text analytics"
+    },
+    "computer_vision_agent": {
+        "categories": ["computer_vision", "profiling"],
+        "description": "Image processing and computer vision tasks"
+    },
+    "business_intelligence_agent": {
+        "categories": ["business_intelligence", "visualization", "profiling"],
+        "description": "Business metrics, KPIs, and strategic insights"
+    },
+    "production_agent": {
+        "categories": ["production", "modeling"],
+        "description": "Model deployment, monitoring, and production operations"
+    },
+    "general_agent": {
+        "categories": ["profiling", "cleaning", "visualization", "code_execution"],
+        "description": "General purpose agent for exploratory analysis"
+    }
+}
+
+# Core tools that should always be available regardless of agent
+CORE_TOOLS = [
+    "profile_dataset",
+    "get_smart_summary",
+    "execute_python_code",
+]
+
+
+def get_tools_for_agent(agent_name: str) -> list:
+    """
+    Get list of tool names relevant to a specific agent.
+    
+    Args:
+        agent_name: Name of the specialist agent
+        
+    Returns:
+        List of tool names the agent can use
+    """
+    if agent_name not in AGENT_TOOL_MAPPING:
+        # Default to general agent tools
+        agent_name = "general_agent"
+    
+    agent_info = AGENT_TOOL_MAPPING[agent_name]
+    categories = agent_info["categories"]
+    
+    # Collect all tools from relevant categories
+    tools = set(CORE_TOOLS)  # Start with core tools
+    
+    for category in categories:
+        if category in TOOL_CATEGORIES:
+            tools.update(TOOL_CATEGORIES[category])
+    
+    return list(tools)
+
+
+def get_tool_categories_for_agent(agent_name: str) -> list:
+    """
+    Get categories of tools relevant to a specific agent.
+    
+    Args:
+        agent_name: Name of the specialist agent
+        
+    Returns:
+        List of tool category names
+    """
+    if agent_name not in AGENT_TOOL_MAPPING:
+        agent_name = "general_agent"
+    
+    return AGENT_TOOL_MAPPING[agent_name]["categories"]
+
+
+def filter_tools_by_names(all_tools: list, tool_names: list) -> list:
+    """
+    Filter tool definitions to only include specified tool names.
+    
+    Args:
+        all_tools: List of all tool definitions (from TOOLS registry)
+        tool_names: List of tool names to include
+        
+    Returns:
+        Filtered list of tool definitions
+    """
+    filtered = []
+    tool_names_set = set(tool_names)
+    
+    for tool in all_tools:
+        if tool.get("type") == "function":
+            function_name = tool.get("function", {}).get("name")
+            if function_name in tool_names_set:
+                # Compress description to reduce token usage
+                compressed_tool = compress_tool_definition(tool)
+                filtered.append(compressed_tool)
+    
+    return filtered
+
+
+def compress_tool_definition(tool: dict) -> dict:
+    """
+    Compress tool definition to reduce token usage.
+    
+    Removes verbose examples and shortens descriptions while keeping
+    essential information for the LLM to use the tool correctly.
+    
+    Args:
+        tool: Tool definition dict
+        
+    Returns:
+        Compressed tool definition
+    """
+    if tool.get("type") != "function":
+        return tool
+    
+    compressed = {
+        "type": "function",
+        "function": {
+            "name": tool["function"]["name"],
+            "description": compress_description(tool["function"]["description"]),
+            "parameters": tool["function"]["parameters"]
+        }
+    }
+    
+    # Compress parameter descriptions
+    if "properties" in compressed["function"]["parameters"]:
+        for param_name, param_info in compressed["function"]["parameters"]["properties"].items():
+            if "description" in param_info:
+                param_info["description"] = compress_description(param_info["description"])
+    
+    return compressed
+
+
+def compress_description(description: str) -> str:
+    """
+    Compress a tool or parameter description.
+    
+    Removes examples, extra whitespace, and verbose explanations
+    while keeping core functionality description.
+    
+    Args:
+        description: Original description
+        
+    Returns:
+        Compressed description
+    """
+    # Remove everything after "Example:" or "Examples:"
+    if "Example:" in description:
+        description = description.split("Example:")[0]
+    if "Examples:" in description:
+        description = description.split("Examples:")[0]
+    
+    # Remove extra whitespace and newlines
+    description = " ".join(description.split())
+    
+    # Truncate if still too long (keep first 150 chars for params, 250 for tools)
+    max_length = 250 if "Use this" in description else 150
+    if len(description) > max_length:
+        description = description[:max_length].rsplit(' ', 1)[0] + "..."
+    
+    return description.strip()
+
+
+def get_agent_description(agent_name: str) -> str:
+    """
+    Get description of what an agent specializes in.
+    
+    Args:
+        agent_name: Name of the specialist agent
+        
+    Returns:
+        Agent description string
+    """
+    if agent_name in AGENT_TOOL_MAPPING:
+        return AGENT_TOOL_MAPPING[agent_name]["description"]
+    return "General purpose data science agent"
+
+
+def suggest_next_agent(current_agent: str, completed_tools: list) -> str:
+    """
+    Suggest the next agent to hand off to based on completed tools.
+    
+    Args:
+        current_agent: Current agent name
+        completed_tools: List of tool names already executed
+        
+    Returns:
+        Suggested next agent name, or None if workflow complete
+    """
+    # Define typical workflow progressions
+    workflows = {
+        "data_quality_agent": "preprocessing_agent",  # After profiling → cleaning
+        "preprocessing_agent": "visualization_agent",   # After cleaning → visualize
+        "visualization_agent": "modeling_agent",        # After EDA → modeling
+        "modeling_agent": "production_agent",           # After training → deploy
+    }
+    
+    # Check if current agent has completed its primary tasks
+    agent_tools = set(get_tools_for_agent(current_agent))
+    completed_set = set(completed_tools)
+    
+    # If less than 30% of agent's tools used, stay with current agent
+    if len(completed_set & agent_tools) / max(len(agent_tools), 1) < 0.3:
+        return current_agent
+    
+    # Suggest next agent in typical workflow
+    return workflows.get(current_agent, None)