agentic-flow 1.8.15 → 1.9.0

package/agentic-llm/training/train.py

@@ -0,0 +1,347 @@
+#!/usr/bin/env python3
+"""
+GPU Agentic Training System for Phi-4 Model
+Optimized for Claude SDK and MCP Tools with Overfitting Prevention
+"""
+
+import os
+import json
+import torch
+import wandb
+from dataclasses import dataclass, field
+from typing import Optional, Dict, List
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    EarlyStoppingCallback,
+    DataCollatorForLanguageModeling
+)
+from datasets import load_dataset, Dataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+import numpy as np
+from pathlib import Path
+import logging
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AgenticTrainingConfig:
+    """Configuration for agentic training system"""
+
+    # Model configuration
+    model_name: str = "microsoft/Phi-4"
+    model_cache_dir: str = "/app/models"
+    output_dir: str = "/app/checkpoints"
+
+    # Training hyperparameters
+    learning_rate: float = 2e-5
+    num_train_epochs: int = 3
+    per_device_train_batch_size: int = 4
+    per_device_eval_batch_size: int = 4
+    gradient_accumulation_steps: int = 4
+    warmup_steps: int = 100
+    max_seq_length: int = 2048
+
+    # LoRA configuration for efficient fine-tuning
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+    lora_target_modules: List[str] = field(default_factory=lambda: ["q_proj", "v_proj", "k_proj", "o_proj"])
+
+    # Overfitting prevention
+    early_stopping_patience: int = 3
+    early_stopping_threshold: float = 0.01
+    eval_steps: int = 100
+    save_steps: int = 200
+    logging_steps: int = 50
+    max_grad_norm: float = 1.0
+    weight_decay: float = 0.01
+    dropout_rate: float = 0.1
+
+    # Validation split
+    validation_split: float = 0.15
+    test_split: float = 0.05
+
+    # MCP/Agentic optimization
+    mcp_tool_focus: bool = True
+    agentic_workflow_focus: bool = True
+
+    # Monitoring
+    use_wandb: bool = True
+    wandb_project: str = "phi4-agentic-optimization"
+
+    # Quantization
+    use_8bit: bool = False
+    use_4bit: bool = True
+
+
+class MCPToolDatasetBuilder:
+    """Build training dataset focused on MCP tool usage and agentic workflows"""
+
+    def __init__(self, config: AgenticTrainingConfig):
+        self.config = config
+
+    def create_mcp_tool_examples(self) -> List[Dict]:
+        """Create synthetic examples for MCP tool calling"""
+        examples = []
+
+        # Tool calling patterns
+        tool_patterns = [
+            {
+                "input": "Search the codebase for authentication functions",
+                "output": "I'll use the Grep tool to search for authentication functions.\n<tool_use>grep(pattern='def.*auth', path='.')</tool_use>"
+            },
+            {
+                "input": "Create a new file for user validation",
+                "output": "I'll create a user validation file.\n<tool_use>write(file_path='./validation/user.py', content='# User validation logic')</tool_use>"
+            },
+            {
+                "input": "Read the configuration file",
+                "output": "I'll read the configuration file.\n<tool_use>read(file_path='./config.json')</tool_use>"
+            },
+            {
+                "input": "List all Python files in the project",
+                "output": "I'll list all Python files.\n<tool_use>glob(pattern='**/*.py')</tool_use>"
+            },
+            {
+                "input": "Execute the test suite",
+                "output": "I'll execute the test suite.\n<tool_use>bash(command='pytest tests/', description='Run test suite')</tool_use>"
+            }
+        ]
+
+        # Agentic workflow patterns
+        workflow_patterns = [
+            {
+                "input": "Implement a new API endpoint with tests",
+                "output": "I'll break this down into steps:\n1. Create endpoint file\n2. Implement logic\n3. Write tests\n4. Run validation\nStarting with step 1..."
+            },
+            {
+                "input": "Debug the authentication error",
+                "output": "Let me investigate:\n1. Check logs\n2. Review auth code\n3. Trace the error\n4. Propose fix\nAnalyzing logs first..."
+            }
+        ]
+
+        examples.extend(tool_patterns)
+        examples.extend(workflow_patterns)
+
+        return examples
+
+    def load_dataset(self) -> Dataset:
+        """Load and prepare training dataset"""
+        # Create MCP-focused examples
+        mcp_examples = self.create_mcp_tool_examples()
+
+        # Convert to dataset format
+        data = {
+            "text": [f"User: {ex['input']}\nAssistant: {ex['output']}" for ex in mcp_examples]
+        }
+
+        dataset = Dataset.from_dict(data)
+        return dataset
+
+
+class AgenticTrainer:
+    """Main training system with overfitting prevention"""
+
+    def __init__(self, config: AgenticTrainingConfig):
+        self.config = config
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {self.device}")
+
+        # Initialize wandb
+        if config.use_wandb:
+            wandb.init(
+                project=config.wandb_project,
+                config=config.__dict__
+            )
+
+    def load_model_and_tokenizer(self):
+        """Load Phi-4 model and tokenizer with quantization"""
+        logger.info(f"Loading model: {self.config.model_name}")
+
+        # Tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.config.model_name,
+            cache_dir=self.config.model_cache_dir,
+            trust_remote_code=True
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+
+        # Model with quantization
+        model = AutoModelForCausalLM.from_pretrained(
+            self.config.model_name,
+            cache_dir=self.config.model_cache_dir,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            load_in_4bit=self.config.use_4bit,
+            load_in_8bit=self.config.use_8bit
+        )
+
+        # Prepare for k-bit training
+        if self.config.use_4bit or self.config.use_8bit:
+            model = prepare_model_for_kbit_training(model)
+
+        # Apply LoRA
+        lora_config = LoraConfig(
+            r=self.config.lora_r,
+            lora_alpha=self.config.lora_alpha,
+            target_modules=self.config.lora_target_modules,
+            lora_dropout=self.config.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+
+        model = get_peft_model(model, lora_config)
+        model.print_trainable_parameters()
+
+        return model, tokenizer
+
+    def prepare_datasets(self, tokenizer):
+        """Prepare training and validation datasets"""
+        logger.info("Preparing datasets...")
+
+        # Build MCP-focused dataset
+        dataset_builder = MCPToolDatasetBuilder(self.config)
+        dataset = dataset_builder.load_dataset()
+
+        # Tokenize
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                max_length=self.config.max_seq_length,
+                padding="max_length"
+            )
+
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=dataset.column_names
+        )
+
+        # Split into train/val/test
+        total_size = len(tokenized_dataset)
+        val_size = int(total_size * self.config.validation_split)
+        test_size = int(total_size * self.config.test_split)
+        train_size = total_size - val_size - test_size
+
+        splits = tokenized_dataset.train_test_split(test_size=val_size + test_size)
+        train_dataset = splits["train"]
+
+        val_test = splits["test"].train_test_split(test_size=test_size)
+        val_dataset = val_test["train"]
+        test_dataset = val_test["test"]
+
+        logger.info(f"Dataset sizes - Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")
+
+        return train_dataset, val_dataset, test_dataset
+
+    def train(self):
+        """Execute training with overfitting prevention"""
+        logger.info("Starting training...")
+
+        # Load model and tokenizer
+        model, tokenizer = self.load_model_and_tokenizer()
+
+        # Prepare datasets
+        train_dataset, val_dataset, test_dataset = self.prepare_datasets(tokenizer)
+
+        # Training arguments with overfitting prevention
+        training_args = TrainingArguments(
+            output_dir=self.config.output_dir,
+            num_train_epochs=self.config.num_train_epochs,
+            per_device_train_batch_size=self.config.per_device_train_batch_size,
+            per_device_eval_batch_size=self.config.per_device_eval_batch_size,
+            gradient_accumulation_steps=self.config.gradient_accumulation_steps,
+            learning_rate=self.config.learning_rate,
+            warmup_steps=self.config.warmup_steps,
+            weight_decay=self.config.weight_decay,
+            max_grad_norm=self.config.max_grad_norm,
+            logging_dir=f"{self.config.output_dir}/logs",
+            logging_steps=self.config.logging_steps,
+            eval_steps=self.config.eval_steps,
+            save_steps=self.config.save_steps,
+            evaluation_strategy="steps",
+            save_strategy="steps",
+            load_best_model_at_end=True,
+            metric_for_best_model="eval_loss",
+            greater_is_better=False,
+            fp16=True,
+            report_to="wandb" if self.config.use_wandb else "none",
+            save_total_limit=3,
+            gradient_checkpointing=True
+        )
+
+        # Data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
+
+        # Early stopping callback
+        early_stopping = EarlyStoppingCallback(
+            early_stopping_patience=self.config.early_stopping_patience,
+            early_stopping_threshold=self.config.early_stopping_threshold
+        )
+
+        # Initialize trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=val_dataset,
+            data_collator=data_collator,
+            callbacks=[early_stopping]
+        )
+
+        # Train
+        trainer.train()
+
+        # Evaluate on test set
+        test_results = trainer.evaluate(test_dataset)
+        logger.info(f"Test results: {test_results}")
+
+        # Save final model
+        trainer.save_model(f"{self.config.output_dir}/final_model")
+        tokenizer.save_pretrained(f"{self.config.output_dir}/final_model")
+
+        # Save metrics
+        metrics = {
+            "test_loss": test_results["eval_loss"],
+            "test_perplexity": np.exp(test_results["eval_loss"]),
+            "config": self.config.__dict__
+        }
+
+        with open(f"{self.config.output_dir}/training_metrics.json", "w") as f:
+            json.dump(metrics, f, indent=2)
+
+        logger.info("Training complete!")
+        return metrics
+
+
+def main():
+    """Main entry point"""
+    # Load configuration
+    config = AgenticTrainingConfig()
+
+    # Initialize trainer
+    trainer = AgenticTrainer(config)
+
+    # Train
+    metrics = trainer.train()
+
+    logger.info("Training pipeline completed successfully!")
+    logger.info(f"Final metrics: {metrics}")
+
+
+if __name__ == "__main__":
+    main()

package/agentic-llm/validation/data/mcp_validation_set.json

@@ -0,0 +1,37 @@
+[
+  {
+    "query": "Read the tsconfig.json file",
+    "expected_tool": "read",
+    "expected_params": {
+      "file_path": "tsconfig.json"
+    }
+  },
+  {
+    "query": "Create a new component file",
+    "expected_tool": "write",
+    "expected_params": {
+      "file_path": "src/components/NewComponent.tsx"
+    }
+  },
+  {
+    "query": "Run the linter",
+    "expected_tool": "bash",
+    "expected_params": {
+      "command": "npm run lint"
+    }
+  },
+  {
+    "query": "Find all error handling code",
+    "expected_tool": "grep",
+    "expected_params": {
+      "pattern": "catch|throw"
+    }
+  },
+  {
+    "query": "List all React components",
+    "expected_tool": "glob",
+    "expected_params": {
+      "pattern": "src/components/**/*.tsx"
+    }
+  }
+]

package/agentic-llm/validation/mcp_validator.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+"""
+MCP Tool Calling Validator
+Validates fine-tuned model on MCP tool accuracy
+"""
+
+import json
+import torch
+from pathlib import Path
+from typing import Dict, List, Tuple
+import re
+import logging
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class MCPToolValidator:
+    """Validate MCP tool calling accuracy"""
+
+    def __init__(self, model_path: str, base_model: str = "microsoft/Phi-4"):
+        self.model_path = model_path
+        self.base_model = base_model
+        self.model = None
+        self.tokenizer = None
+
+        self._load_model()
+
+    def _load_model(self):
+        """Load fine-tuned model"""
+        logger.info(f"Loading model from {self.model_path}")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_path,
+            trust_remote_code=True
+        )
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        # Check if this is a PEFT model
+        peft_config_path = Path(self.model_path) / "adapter_config.json"
+
+        if peft_config_path.exists():
+            # Load base model first
+            base = AutoModelForCausalLM.from_pretrained(
+                self.base_model,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True
+            )
+            # Load PEFT adapter
+            self.model = PeftModel.from_pretrained(base, self.model_path)
+            self.model = self.model.merge_and_unload()  # Merge for inference
+        else:
+            # Load full model
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True
+            )
+
+        self.model.eval()
+        logger.info("Model loaded successfully")
+
+    def parse_tool_call(self, response: str) -> Dict:
+        """Parse tool call from response"""
+        # Look for <tool_use>...</tool_use> pattern
+        pattern = r'<tool_use>(.*?)</tool_use>'
+        match = re.search(pattern, response, re.DOTALL)
+
+        if not match:
+            return {"found": False}
+
+        tool_call_str = match.group(1).strip()
+
+        # Extract tool name and parameters
+        # Format: tool_name(param1='value1', param2='value2')
+        tool_match = re.match(r'(\w+)\((.*)\)', tool_call_str)
+
+        if not tool_match:
+            return {"found": True, "parsed": False}
+
+        tool_name = tool_match.group(1)
+        params_str = tool_match.group(2)
+
+        # Parse parameters
+        params = {}
+        if params_str:
+            # Simple parameter parsing
+            param_pattern = r"(\w+)=(['\"])(.*?)\2"
+            for param_match in re.finditer(param_pattern, params_str):
+                key = param_match.group(1)
+                value = param_match.group(3)
+                params[key] = value
+
+        return {
+            "found": True,
+            "parsed": True,
+            "tool_name": tool_name,
+            "parameters": params
+        }
+
+    def generate_response(self, prompt: str, max_tokens: int = 150) -> str:
+        """Generate response from model"""
+        formatted_prompt = f"User: {prompt}\nAssistant:"
+
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=0.1,
+                do_sample=True,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract assistant response
+        if "Assistant:" in response:
+            response = response.split("Assistant:")[-1].strip()
+
+        return response
+
+    def validate_tool_accuracy(self, test_cases: List[Dict]) -> Dict:
+        """Validate tool calling accuracy on test cases"""
+        logger.info(f"Validating on {len(test_cases)} test cases...")
+
+        results = {
+            "total": len(test_cases),
+            "correct_tool": 0,
+            "correct_params": 0,
+            "tool_found": 0,
+            "tool_parsed": 0,
+            "accuracy": 0.0,
+            "param_accuracy": 0.0,
+            "details": []
+        }
+
+        for i, test_case in enumerate(test_cases):
+            query = test_case["query"]
+            expected_tool = test_case["expected_tool"]
+            expected_params = test_case.get("expected_params", {})
+
+            # Generate response
+            response = self.generate_response(query)
+
+            # Parse tool call
+            parsed = self.parse_tool_call(response)
+
+            # Check correctness
+            tool_found = parsed.get("found", False)
+            tool_parsed = parsed.get("parsed", False)
+            correct_tool = tool_parsed and parsed.get("tool_name") == expected_tool
+            correct_params = False
+
+            if correct_tool and expected_params:
+                # Check if all expected params are present
+                actual_params = parsed.get("parameters", {})
+                correct_params = all(
+                    key in actual_params for key in expected_params.keys()
+                )
+
+            # Update counts
+            if tool_found:
+                results["tool_found"] += 1
+            if tool_parsed:
+                results["tool_parsed"] += 1
+            if correct_tool:
+                results["correct_tool"] += 1
+            if correct_params:
+                results["correct_params"] += 1
+
+            # Record details
+            results["details"].append({
+                "query": query,
+                "expected_tool": expected_tool,
+                "response": response,
+                "parsed": parsed,
+                "correct_tool": correct_tool,
+                "correct_params": correct_params
+            })
+
+            if (i + 1) % 10 == 0:
+                logger.info(f"Processed {i + 1}/{len(test_cases)} test cases")
+
+        # Calculate final metrics
+        results["accuracy"] = (results["correct_tool"] / results["total"]) * 100
+        results["param_accuracy"] = (results["correct_params"] / results["total"]) * 100
+        results["parse_rate"] = (results["tool_parsed"] / results["tool_found"]) * 100 if results["tool_found"] > 0 else 0
+
+        logger.info("\nValidation Results:")
+        logger.info(f"  Tool Detection: {results['tool_found']}/{results['total']} ({results['tool_found']/results['total']*100:.1f}%)")
+        logger.info(f"  Tool Parsing: {results['tool_parsed']}/{results['tool_found']} ({results['parse_rate']:.1f}%)")
+        logger.info(f"  Tool Accuracy: {results['correct_tool']}/{results['total']} ({results['accuracy']:.1f}%)")
+        logger.info(f"  Param Accuracy: {results['correct_params']}/{results['total']} ({results['param_accuracy']:.1f}%)")
+
+        return results
+
+
+def main():
+    """Run validation"""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Validate MCP tool calling")
+    parser.add_argument("--model-path", type=str, required=True, help="Path to fine-tuned model")
+    parser.add_argument("--base-model", type=str, default="microsoft/Phi-4", help="Base model name")
+    parser.add_argument("--output-dir", type=str, default="/app/validation/results", help="Output directory")
+
+    args = parser.parse_args()
+
+    # Create validator
+    validator = MCPToolValidator(
+        model_path=args.model_path,
+        base_model=args.base_model
+    )
+
+    # Load test cases
+    test_cases = [
+        {"query": "Read the package.json file", "expected_tool": "read", "expected_params": {"file_path": "package.json"}},
+        {"query": "Create a new utils file", "expected_tool": "write", "expected_params": {"file_path": "utils.ts"}},
+        {"query": "Run the tests", "expected_tool": "bash", "expected_params": {"command": "npm test"}},
+        {"query": "Search for TODO comments", "expected_tool": "grep", "expected_params": {"pattern": "TODO"}},
+        {"query": "List all TypeScript files", "expected_tool": "glob", "expected_params": {"pattern": "**/*.ts"}},
+        {"query": "Update the port to 8080", "expected_tool": "edit", "expected_params": {"old_string": "3000", "new_string": "8080"}},
+        {"query": "Show me the config file", "expected_tool": "read", "expected_params": {"file_path": "config"}},
+        {"query": "Install dependencies", "expected_tool": "bash", "expected_params": {"command": "npm install"}},
+        {"query": "Find all API endpoints", "expected_tool": "grep", "expected_params": {"pattern": "api"}},
+        {"query": "Show all JavaScript files", "expected_tool": "glob", "expected_params": {"pattern": "**/*.js"}},
+    ]
+
+    # Validate
+    results = validator.validate_tool_accuracy(test_cases)
+
+    # Save results
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    output_file = output_dir / "mcp_validation_results.json"
+    with open(output_file, "w") as f:
+        json.dump(results, f, indent=2)
+
+    logger.info(f"\nResults saved to {output_file}")
+
+    # Print summary
+    logger.info("\n" + "=" * 60)
+    logger.info("VALIDATION SUMMARY")
+    logger.info("=" * 60)
+    logger.info(f"Overall Tool Accuracy: {results['accuracy']:.1f}%")
+    logger.info(f"Parameter Accuracy: {results['param_accuracy']:.1f}%")
+
+    if results['accuracy'] >= 85:
+        logger.info("✅ EXCELLENT - Model is ready for deployment")
+    elif results['accuracy'] >= 75:
+        logger.info("✓ GOOD - Model performs well")
+    elif results['accuracy'] >= 65:
+        logger.info("⚠ ACCEPTABLE - Consider more training")
+    else:
+        logger.info("❌ NEEDS IMPROVEMENT - More training required")
+
+
+if __name__ == "__main__":
+    main()