For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (
- [ ]) syntax for tracking.
Goal: Build a multi-provider ML framework that replicates ml-intern functionality with support for Claude, OpenAI, DeepSeek, and Mistral, enabling automated ML workflows (dataset collection → training → deployment).
Architecture: A modular Python framework with pluggable LLM providers, workflow engine, and CLI interface. Core abstracts provider differences; workflows compose reusable steps; plugins enable extensibility.
Tech Stack: Python 3.11+, Typer (CLI), Pydantic (config), Structlog (logging), Hugging Face Transformers, pytest
- Python minimum: 3.11
- License: Apache 2.0
- Providers in v1.0: Claude, OpenAI, DeepSeek, Mistral (others via plugins)
- Output directory: ~/ml-agent (respects XDG paths)
- Config format: YAML + .env
- Logging: Structured JSON logs to ~/.ml-agent/logs/
- Test coverage minimum: 80% excluding examples
- Public release: p7dotorg organization on GitHub
Files:
- Create:
pyproject.toml - Create:
src/ml_agent/__init__.py - Create:
src/ml_agent/core/__init__.py - Create:
.github/workflows/test.yml
Interfaces:
-
Produces: Package structure ready for imports (
from ml_agent.core import ...) -
Step 1: Create pyproject.toml with all dependencies
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "ml-agent"
version = "1.0.0"
description = "Multi-provider ML workflow framework"
readme = "README.md"
license = {text = "Apache-2.0"}
authors = [{name = "p7dotorg"}]
requires-python = ">=3.11"
dependencies = [
"typer[all]==0.9.0",
"pydantic==2.5.0",
"pydantic-settings==2.1.0",
"structlog==23.2.0",
"python-dotenv==1.0.0",
"anthropic>=0.18.0",
"openai>=1.3.0",
"httpx>=0.25.0",
"PyYAML==6.0",
"requests==2.31.0",
"tqdm==4.66.1",
"pandas==2.1.1",
]
[project.optional-dependencies]
dev = [
"pytest>=7.4.0",
"pytest-cov>=4.1.0",
"pytest-asyncio>=0.21.0",
"black==23.11.0",
"isort==5.12.0",
"mypy==1.7.0",
"ruff==0.1.8",
]
[project.scripts]
ml-agent = "ml_agent.cli:app"
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = "test_*.py"
python_classes = "Test*"
python_functions = "test_*"- Step 2: Create directory structure
mkdir -p src/ml_agent/{core,providers,auth,workflows,utils,plugins,cli}
mkdir -p tests/{unit,integration,fixtures}
mkdir -p examples/{latex-explainer,simple-workflow}
mkdir -p docs- Step 3: Create init.py files
# src/ml_agent/__init__.py
"""ML Agent Framework - Multi-provider ML workflow automation."""
__version__ = "1.0.0"
__author__ = "p7dotorg"
# src/ml_agent/core/__init__.py
"""Core components of ML Agent Framework."""- Step 4: Create GitHub Actions test workflow
# .github/workflows/test.yml
name: Tests
on: [push, pull_request]
jobs:
test:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.11", "3.12"]
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
python -m pip install -e ".[dev]"
- name: Run tests with coverage
run: |
pytest --cov=src/ml_agent --cov-report=xml tests/
- name: Upload coverage
uses: codecov/codecov-action@v3- Step 5: Create .gitignore
__pycache__/
*.py[cod]
*$py.class
.env
.env.local
auth.json
*.egg-info/
dist/
build/
.pytest_cache/
.mypy_cache/
.coverage
htmlcov/
logs/
*.log
.DS_Store
- Step 6: Commit
git add pyproject.toml src/ tests/ .github/ .gitignore
git commit -m "feat: setup project structure and dependencies"Files:
- Create:
src/ml_agent/core/config.py - Create:
src/ml_agent/core/exceptions.py - Create:
tests/unit/test_config.py
Interfaces:
-
Produces:
Configclass (Pydantic model), custom exceptions -
Step 1: Create config.py
# src/ml_agent/core/config.py
from pathlib import Path
from typing import Optional
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings
class ProviderConfig(BaseModel):
"""Configuration for a specific LLM provider."""
api_key: Optional[str] = Field(None, description="API key for provider")
base_url: Optional[str] = Field(None, description="Custom base URL")
timeout: int = Field(30, description="Request timeout in seconds")
max_retries: int = Field(3, description="Maximum retry attempts")
class WorkflowConfig(BaseModel):
"""Configuration for a workflow execution."""
name: str
provider: str
steps: dict = Field(default_factory=dict)
max_iterations: int = 300
timeout_minutes: int = 120
class Config(BaseSettings):
"""Main configuration for ML Agent."""
# Directories
home_dir: Path = Field(default_factory=lambda: Path.home() / ".ml-agent")
log_dir: Path = Field(default_factory=lambda: Path.home() / ".ml-agent" / "logs")
config_file: Path = Field(default_factory=lambda: Path.home() / ".ml-agent" / "config.yaml")
auth_file: Path = Field(default_factory=lambda: Path.home() / ".ml-agent" / "auth.json")
# Logging
log_level: str = Field("INFO", description="Log level")
log_format: str = Field("json", description="json or text")
# Providers
providers: dict[str, ProviderConfig] = Field(default_factory=dict)
# General
debug: bool = False
class Config:
env_prefix = "ML_AGENT_"
yaml_file = None
def ensure_directories(self):
"""Create necessary directories."""
self.home_dir.mkdir(parents=True, exist_ok=True)
self.log_dir.mkdir(parents=True, exist_ok=True)
@classmethod
def from_file(cls, path: Path) -> "Config":
"""Load config from YAML file."""
import yaml
with open(path) as f:
data = yaml.safe_load(f) or {}
return cls(**data)- Step 2: Create exceptions.py
# src/ml_agent/core/exceptions.py
"""Custom exceptions for ML Agent Framework."""
class MLAgentException(Exception):
"""Base exception for ML Agent."""
pass
class ProviderException(MLAgentException):
"""Provider-related errors."""
pass
class AuthenticationError(ProviderException):
"""Authentication/authorization errors."""
pass
class RateLimitError(ProviderException):
"""Rate limit exceeded."""
pass
class WorkflowException(MLAgentException):
"""Workflow execution errors."""
pass
class ConfigurationError(MLAgentException):
"""Configuration errors."""
pass
class ValidationError(MLAgentException):
"""Data validation errors."""
pass- Step 3: Create test_config.py
# tests/unit/test_config.py
import pytest
from pathlib import Path
from ml_agent.core.config import Config, ProviderConfig
def test_config_defaults():
"""Test default configuration values."""
config = Config()
assert config.log_level == "INFO"
assert config.debug is False
assert isinstance(config.home_dir, Path)
def test_provider_config():
"""Test provider configuration."""
provider = ProviderConfig(api_key="test-key")
assert provider.api_key == "test-key"
assert provider.timeout == 30
assert provider.max_retries == 3
def test_ensure_directories(tmp_path):
"""Test directory creation."""
config = Config(home_dir=tmp_path / ".ml-agent")
config.ensure_directories()
assert config.home_dir.exists()
assert config.log_dir.exists()- Step 4: Run tests
pytest tests/unit/test_config.py -v
# Expected: PASS- Step 5: Commit
git add src/ml_agent/core/{config,exceptions}.py tests/unit/test_config.py
git commit -m "feat: add configuration management and exceptions"Files:
- Create:
src/ml_agent/core/logger.py - Create:
tests/unit/test_logger.py
Interfaces:
-
Consumes:
Config(from Task 2) -
Produces:
setup_logging()function,get_logger()helper -
Step 1: Create logger.py
# src/ml_agent/core/logger.py
import logging
import structlog
from pathlib import Path
from typing import Optional
def setup_logging(
level: str = "INFO",
log_dir: Path = None,
format: str = "json"
):
"""Configure structlog and standard logging."""
# Standard library config
logging.basicConfig(
format="%(message)s",
level=getattr(logging, level),
)
# structlog config
structlog.configure(
processors=[
structlog.stdlib.filter_by_level,
structlog.stdlib.add_logger_name,
structlog.stdlib.add_log_level,
structlog.stdlib.PositionalArgumentsFormatter(),
structlog.processors.TimeStamper(fmt="iso"),
structlog.processors.StackInfoRenderer(),
structlog.processors.format_exc_info,
structlog.processors.UnicodeDecoder(),
structlog.processors.JSONRenderer()
if format == "json"
else structlog.dev.ConsoleRenderer(),
],
context_class=dict,
logger_factory=structlog.stdlib.LoggerFactory(),
cache_logger_on_first_use=True,
)
# File handler if log_dir provided
if log_dir:
log_dir.mkdir(parents=True, exist_ok=True)
file_handler = logging.FileHandler(
log_dir / "ml-agent.log"
)
file_handler.setFormatter(
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)
logging.getLogger().addHandler(file_handler)
def get_logger(name: str = __name__):
"""Get configured logger instance."""
return structlog.get_logger(name)- Step 2: Create test_logger.py
# tests/unit/test_logger.py
import logging
from ml_agent.core.logger import setup_logging, get_logger
def test_setup_logging():
"""Test logging setup."""
setup_logging(level="DEBUG")
logger = get_logger("test")
assert logger is not None
def test_get_logger():
"""Test getting logger instance."""
setup_logging()
logger = get_logger("test_module")
assert logger is not None
# Should not raise
logger.info("test message")- Step 3: Run tests
pytest tests/unit/test_logger.py -v
# Expected: PASS- Step 4: Commit
git add src/ml_agent/core/logger.py tests/unit/test_logger.py
git commit -m "feat: add structured logging with structlog"Files:
- Create:
src/ml_agent/providers/base.py - Create:
tests/unit/test_providers_base.py
Interfaces:
-
Consumes:
Config, logging -
Produces:
BaseProviderabstract class -
Step 1: Create base.py
# src/ml_agent/providers/base.py
from abc import ABC, abstractmethod
from typing import Optional
from dataclasses import dataclass
import structlog
@dataclass
class Message:
"""LLM message format."""
role: str # "user", "assistant", "system"
content: str
class BaseProvider(ABC):
"""Abstract base for LLM providers."""
name: str # Must be defined by subclass
def __init__(self, api_key: str, **kwargs):
"""Initialize provider."""
self.api_key = api_key
self.logger = structlog.get_logger(self.__class__.__name__)
self.config = kwargs
@abstractmethod
async def complete(
self,
messages: list[Message],
max_tokens: int = 2048,
temperature: float = 0.7,
**kwargs
) -> str:
"""Generate completion from messages."""
pass
@abstractmethod
def validate_credentials(self) -> bool:
"""Validate API credentials."""
pass
def __repr__(self) -> str:
return f"{self.__class__.__name__}(api_key=***)"- Step 2: Create test_providers_base.py
# tests/unit/test_providers_base.py
import pytest
from ml_agent.providers.base import BaseProvider, Message
class MockProvider(BaseProvider):
name = "mock"
async def complete(self, messages, max_tokens=2048, temperature=0.7, **kwargs):
return "Mock response"
def validate_credentials(self):
return True
def test_base_provider_initialization():
"""Test provider initialization."""
provider = MockProvider(api_key="test-key")
assert provider.api_key == "test-key"
assert provider.name == "mock"
def test_message_creation():
"""Test message creation."""
msg = Message(role="user", content="Hello")
assert msg.role == "user"
assert msg.content == "Hello"- Step 3: Run tests
pytest tests/unit/test_providers_base.py -v
# Expected: PASS- Step 4: Commit
git add src/ml_agent/providers/base.py tests/unit/test_providers_base.py
git commit -m "feat: add base provider abstraction"Files:
- Create:
src/ml_agent/providers/claude.py - Create:
tests/unit/test_providers_claude.py
Interfaces:
-
Consumes:
BaseProvider(Task 4), Anthropic SDK -
Produces:
ClaudeProviderclass -
Step 1: Create claude.py
# src/ml_agent/providers/claude.py
import anthropic
from typing import Optional
from ml_agent.providers.base import BaseProvider, Message
from ml_agent.core.exceptions import AuthenticationError, ProviderException
class ClaudeProvider(BaseProvider):
"""Anthropic Claude provider."""
name = "claude"
default_model = "claude-3-5-sonnet-20241022"
def __init__(self, api_key: str, model: Optional[str] = None, **kwargs):
super().__init__(api_key, **kwargs)
self.model = model or self.default_model
self.client = anthropic.Anthropic(api_key=api_key)
async def complete(
self,
messages: list[Message],
max_tokens: int = 2048,
temperature: float = 0.7,
**kwargs
) -> str:
"""Generate completion using Claude."""
try:
formatted_messages = [
{"role": msg.role, "content": msg.content}
for msg in messages
]
response = self.client.messages.create(
model=self.model,
max_tokens=max_tokens,
temperature=temperature,
messages=formatted_messages,
**kwargs
)
return response.content[0].text
except anthropic.AuthenticationError as e:
raise AuthenticationError(f"Claude authentication failed: {e}")
except anthropic.RateLimitError as e:
raise ProviderException(f"Claude rate limited: {e}")
except Exception as e:
raise ProviderException(f"Claude error: {e}")
def validate_credentials(self) -> bool:
"""Validate API key."""
try:
# Try to make a minimal request
self.client.messages.create(
model=self.model,
max_tokens=10,
messages=[{"role": "user", "content": "test"}]
)
return True
except:
return False- Step 2: Create test_providers_claude.py
# tests/unit/test_providers_claude.py
import pytest
from unittest.mock import Mock, AsyncMock, patch
from ml_agent.providers.claude import ClaudeProvider
from ml_agent.providers.base import Message
@pytest.fixture
def claude_provider():
with patch('anthropic.Anthropic'):
return ClaudeProvider(api_key="sk-ant-test")
def test_claude_initialization(claude_provider):
"""Test Claude provider initialization."""
assert claude_provider.name == "claude"
assert claude_provider.api_key == "sk-ant-test"
assert claude_provider.model == ClaudeProvider.default_model
def test_claude_repr(claude_provider):
"""Test provider string representation."""
assert "***" in repr(claude_provider)
assert "claude" in repr(claude_provider).lower()- Step 3: Run tests
pytest tests/unit/test_providers_claude.py -v
# Expected: PASS- Step 4: Commit
git add src/ml_agent/providers/claude.py tests/unit/test_providers_claude.py
git commit -m "feat: add Claude provider implementation"Files:
- Create:
src/ml_agent/providers/openai.py - Create:
tests/unit/test_providers_openai.py
Interfaces:
-
Consumes:
BaseProvider(Task 4), OpenAI SDK -
Produces:
OpenAIProviderclass -
Step 1: Create openai.py
# src/ml_agent/providers/openai.py
import openai
from typing import Optional
from ml_agent.providers.base import BaseProvider, Message
from ml_agent.core.exceptions import AuthenticationError, ProviderException
class OpenAIProvider(BaseProvider):
"""OpenAI GPT provider."""
name = "openai"
default_model = "gpt-4-turbo"
def __init__(self, api_key: str, model: Optional[str] = None, **kwargs):
super().__init__(api_key, **kwargs)
self.model = model or self.default_model
self.client = openai.OpenAI(api_key=api_key)
async def complete(
self,
messages: list[Message],
max_tokens: int = 2048,
temperature: float = 0.7,
**kwargs
) -> str:
"""Generate completion using OpenAI."""
try:
formatted_messages = [
{"role": msg.role, "content": msg.content}
for msg in messages
]
response = self.client.chat.completions.create(
model=self.model,
max_tokens=max_tokens,
temperature=temperature,
messages=formatted_messages,
**kwargs
)
return response.choices[0].message.content
except openai.AuthenticationError as e:
raise AuthenticationError(f"OpenAI authentication failed: {e}")
except openai.RateLimitError as e:
raise ProviderException(f"OpenAI rate limited: {e}")
except Exception as e:
raise ProviderException(f"OpenAI error: {e}")
def validate_credentials(self) -> bool:
"""Validate API key."""
try:
self.client.chat.completions.create(
model=self.model,
max_tokens=10,
messages=[{"role": "user", "content": "test"}]
)
return True
except:
return False- Step 2: Create test_providers_openai.py
# tests/unit/test_providers_openai.py
import pytest
from unittest.mock import patch
from ml_agent.providers.openai import OpenAIProvider
@pytest.fixture
def openai_provider():
with patch('openai.OpenAI'):
return OpenAIProvider(api_key="sk-test")
def test_openai_initialization(openai_provider):
"""Test OpenAI provider initialization."""
assert openai_provider.name == "openai"
assert openai_provider.model == OpenAIProvider.default_model- Step 3: Run tests
pytest tests/unit/test_providers_openai.py -v
# Expected: PASS- Step 4: Commit
git add src/ml_agent/providers/openai.py tests/unit/test_providers_openai.py
git commit -m "feat: add OpenAI provider implementation"Files:
- Create:
src/ml_agent/providers/registry.py - Create:
src/ml_agent/providers/__init__.py - Create:
tests/unit/test_providers_registry.py
Interfaces:
-
Consumes: All provider implementations (Tasks 5-6)
-
Produces:
ProviderRegistryclass,get_provider()factory function -
Step 1: Create registry.py
# src/ml_agent/providers/registry.py
from typing import Dict, Type
from ml_agent.providers.base import BaseProvider
from ml_agent.providers.claude import ClaudeProvider
from ml_agent.providers.openai import OpenAIProvider
class ProviderRegistry:
"""Registry for available LLM providers."""
_providers: Dict[str, Type[BaseProvider]] = {
"claude": ClaudeProvider,
"openai": OpenAIProvider,
}
@classmethod
def register(cls, name: str, provider_class: Type[BaseProvider]):
"""Register a new provider."""
cls._providers[name] = provider_class
@classmethod
def get(cls, name: str, api_key: str, **kwargs) -> BaseProvider:
"""Get provider instance."""
if name not in cls._providers:
raise ValueError(f"Unknown provider: {name}")
return cls._providers[name](api_key=api_key, **kwargs)
@classmethod
def list_available(cls) -> list[str]:
"""List all available providers."""
return list(cls._providers.keys())
def get_provider(name: str, api_key: str, **kwargs) -> BaseProvider:
"""Factory function to get provider instance."""
return ProviderRegistry.get(name, api_key, **kwargs)- Step 2: Create providers/init.py
# src/ml_agent/providers/__init__.py
from ml_agent.providers.base import BaseProvider, Message
from ml_agent.providers.registry import ProviderRegistry, get_provider
__all__ = ["BaseProvider", "Message", "ProviderRegistry", "get_provider"]- Step 3: Create test_providers_registry.py
# tests/unit/test_providers_registry.py
import pytest
from ml_agent.providers.registry import ProviderRegistry, get_provider
def test_list_available_providers():
"""Test listing available providers."""
providers = ProviderRegistry.list_available()
assert "claude" in providers
assert "openai" in providers
def test_get_unknown_provider():
"""Test error on unknown provider."""
with pytest.raises(ValueError, match="Unknown provider"):
ProviderRegistry.get("unknown", "key")
def test_factory_function():
"""Test factory function."""
# Should not raise (won't actually connect)
providers = ProviderRegistry.list_available()
assert len(providers) > 0- Step 4: Run tests
pytest tests/unit/test_providers_*.py -v
# Expected: ALL PASS- Step 5: Commit
git add src/ml_agent/providers/{registry,__init__}.py tests/unit/test_providers_registry.py
git commit -m "feat: add provider registry and factory"Files:
- Create:
src/ml_agent/auth/manager.py - Create:
src/ml_agent/auth/strategies.py - Create:
tests/unit/test_auth.py
Interfaces:
-
Consumes:
Config, exceptions -
Produces:
AuthManagerclass, credential resolution -
Step 1: Create strategies.py
# src/ml_agent/auth/strategies.py
from abc import ABC, abstractmethod
from pathlib import Path
import os
import json
from typing import Optional
class AuthStrategy(ABC):
"""Base authentication strategy."""
@abstractmethod
def get_credentials(self, provider: str) -> Optional[str]:
"""Get credentials for provider."""
pass
class EnvVarStrategy(AuthStrategy):
"""Get credentials from environment variables."""
def get_credentials(self, provider: str) -> Optional[str]:
"""Look for {PROVIDER}_API_KEY or similar."""
env_keys = [
f"{provider.upper()}_API_KEY",
f"{provider.upper().replace('-', '_')}_API_KEY",
]
for key in env_keys:
if value := os.getenv(key):
return value
return None
class FileStrategy(AuthStrategy):
"""Get credentials from auth.json file."""
def __init__(self, auth_file: Path):
self.auth_file = auth_file
def get_credentials(self, provider: str) -> Optional[str]:
"""Read from auth.json."""
if not self.auth_file.exists():
return None
try:
with open(self.auth_file) as f:
data = json.load(f)
return data.get("providers", {}).get(provider, {}).get("api_key")
except (json.JSONDecodeError, KeyError):
return None
class CLIArgStrategy(AuthStrategy):
"""Get credentials from CLI arguments."""
def __init__(self, api_key: Optional[str] = None):
self.api_key = api_key
def get_credentials(self, provider: str) -> Optional[str]:
"""Return CLI-provided key if matches provider."""
return self.api_key
class InteractiveStrategy(AuthStrategy):
"""Prompt user for credentials."""
def __init__(self, save_to_file: bool = False, auth_file: Path = None):
self.save_to_file = save_to_file
self.auth_file = auth_file
def get_credentials(self, provider: str) -> Optional[str]:
"""Prompt user for API key."""
import getpass
key = getpass.getpass(f"Enter {provider.upper()} API key: ")
if self.save_to_file and self.auth_file and key:
self._save_to_file(provider, key)
return key if key else None
def _save_to_file(self, provider: str, key: str):
"""Save credentials to auth.json."""
self.auth_file.parent.mkdir(parents=True, exist_ok=True)
data = {}
if self.auth_file.exists():
with open(self.auth_file) as f:
data = json.load(f)
if "providers" not in data:
data["providers"] = {}
data["providers"][provider] = {"api_key": key}
with open(self.auth_file, "w") as f:
json.dump(data, f, indent=2)- Step 2: Create manager.py
# src/ml_agent/auth/manager.py
from pathlib import Path
from typing import Optional
from ml_agent.auth.strategies import (
AuthStrategy, EnvVarStrategy, FileStrategy,
CLIArgStrategy, InteractiveStrategy
)
from ml_agent.core.exceptions import AuthenticationError
class AuthManager:
"""Manages authentication across providers."""
def __init__(self, auth_file: Path = None, cli_api_key: Optional[str] = None):
self.auth_file = auth_file or Path.home() / ".ml-agent" / "auth.json"
self.strategies = [
CLIArgStrategy(cli_api_key),
FileStrategy(self.auth_file),
EnvVarStrategy(),
InteractiveStrategy(save_to_file=True, auth_file=self.auth_file),
]
def get_api_key(self, provider: str) -> str:
"""Get API key for provider using resolution order."""
for strategy in self.strategies:
if key := strategy.get_credentials(provider):
return key
raise AuthenticationError(
f"No credentials found for provider '{provider}'. "
f"Please provide via CLI, env var, or auth.json"
)
def validate_provider_auth(self, provider_instance) -> bool:
"""Validate provider credentials."""
return provider_instance.validate_credentials()- Step 3: Create test_auth.py
# tests/unit/test_auth.py
import pytest
from pathlib import Path
from ml_agent.auth.manager import AuthManager
from ml_agent.auth.strategies import (
EnvVarStrategy, FileStrategy, CLIArgStrategy
)
from ml_agent.core.exceptions import AuthenticationError
def test_cli_arg_strategy():
"""Test CLI argument strategy."""
strategy = CLIArgStrategy(api_key="test-key")
assert strategy.get_credentials("claude") == "test-key"
def test_env_var_strategy(monkeypatch):
"""Test environment variable strategy."""
monkeypatch.setenv("CLAUDE_API_KEY", "env-key")
strategy = EnvVarStrategy()
assert strategy.get_credentials("claude") == "env-key"
def test_file_strategy(tmp_path):
"""Test file-based strategy."""
auth_file = tmp_path / "auth.json"
strategy = FileStrategy(auth_file)
assert strategy.get_credentials("claude") is None
def test_auth_manager_get_api_key_cli():
"""Test AuthManager with CLI key."""
manager = AuthManager(cli_api_key="cli-key")
assert manager.get_api_key("claude") == "cli-key"
def test_auth_manager_no_credentials():
"""Test error when no credentials found."""
manager = AuthManager()
with pytest.raises(AuthenticationError):
manager.get_api_key("nonexistent")- Step 4: Run tests
pytest tests/unit/test_auth.py -v
# Expected: PASS (except interactive test)- Step 5: Create init.py
# src/ml_agent/auth/__init__.py
from ml_agent.auth.manager import AuthManager
from ml_agent.auth.strategies import AuthStrategy
__all__ = ["AuthManager", "AuthStrategy"]- Step 6: Commit
git add src/ml_agent/auth/ tests/unit/test_auth.py
git commit -m "feat: add authentication manager with multiple strategies"Files:
- Create:
src/ml_agent/cli.py - Create:
tests/unit/test_cli.py
Interfaces:
-
Consumes:
AuthManager,ProviderRegistry,Config -
Produces: Typer CLI app with commands
-
Step 1: Create cli.py
# src/ml_agent/cli.py
import typer
from pathlib import Path
from typing import Optional
import structlog
from ml_agent.core.config import Config
from ml_agent.core.logger import setup_logging
from ml_agent.auth.manager import AuthManager
from ml_agent.providers.registry import ProviderRegistry
app = typer.Typer(
name="ml-agent",
help="Multi-provider ML workflow framework"
)
def get_config() -> Config:
"""Get or create configuration."""
config = Config()
config.ensure_directories()
return config
@app.command()
def list_providers():
"""List available LLM providers."""
config = get_config()
setup_logging(config.log_level, config.log_dir)
logger = structlog.get_logger(__name__)
providers = ProviderRegistry.list_available()
typer.echo("Available providers:")
for provider in providers:
typer.echo(f" • {provider}")
logger.info("Listed providers", count=len(providers))
@app.command()
def validate_auth(
provider: str = typer.Option(..., help="Provider name"),
api_key: Optional[str] = typer.Option(None, help="API key (or use env/auth.json)"),
):
"""Validate credentials for a provider."""
config = get_config()
setup_logging(config.log_level, config.log_dir)
logger = structlog.get_logger(__name__)
auth_manager = AuthManager(cli_api_key=api_key)
try:
key = auth_manager.get_api_key(provider)
provider_instance = ProviderRegistry.get(provider, key)
if provider_instance.validate_credentials():
typer.secho(f"✓ {provider} credentials valid!", fg=typer.colors.GREEN)
logger.info("Auth valid", provider=provider)
else:
typer.secho(f"✗ {provider} credentials invalid!", fg=typer.colors.RED)
logger.warning("Auth invalid", provider=provider)
except Exception as e:
typer.secho(f"✗ Error: {e}", fg=typer.colors.RED)
logger.error("Auth error", provider=provider, error=str(e))
raise typer.Exit(1)
@app.command()
def version():
"""Show version."""
from ml_agent import __version__
typer.echo(f"ml-agent {__version__}")
@app.callback()
def main(
debug: bool = typer.Option(False, "--debug", help="Enable debug logging"),
):
"""ML Agent Framework - Multi-provider ML workflows."""
if debug:
import sys
sys.exit(0) # Placeholder for debug setup
if __name__ == "__main__":
app()- Step 2: Create test_cli.py
# tests/unit/test_cli.py
from typer.testing import CliRunner
from ml_agent.cli import app
runner = CliRunner()
def test_list_providers():
"""Test list-providers command."""
result = runner.invoke(app, ["list-providers"])
assert result.exit_code == 0
assert "claude" in result.stdout.lower()
assert "openai" in result.stdout.lower()
def test_version():
"""Test version command."""
result = runner.invoke(app, ["version"])
assert result.exit_code == 0
assert "ml-agent" in result.stdout
def test_help():
"""Test help command."""
result = runner.invoke(app, ["--help"])
assert result.exit_code == 0
assert "Multi-provider" in result.stdout- Step 3: Run tests
pytest tests/unit/test_cli.py -v
# Expected: PASS- Step 4: Update pyproject.toml
Já foi feito na Task 1, mas confirmar:
[project.scripts]
ml-agent = "ml_agent.cli:app"- Step 5: Commit
git add src/ml_agent/cli.py tests/unit/test_cli.py
git commit -m "feat: add CLI interface with Typer"Files:
- Create:
src/ml_agent/workflows/base.py - Create:
src/ml_agent/workflows/registry.py - Create:
src/ml_agent/workflows/__init__.py - Create:
tests/unit/test_workflows.py
Interfaces:
-
Consumes: Providers, config
-
Produces:
Workflowbase class,WorkflowRegistry -
Step 1: Create base.py
# src/ml_agent/workflows/base.py
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
from dataclasses import dataclass
import structlog
@dataclass
class WorkflowStep:
"""Single step in a workflow."""
name: str
task: str # Task for LLM
validation: Optional[callable] = None
class Workflow(ABC):
"""Base class for all workflows."""
name: str
description: str
def __init__(self, provider, config: Dict[str, Any] = None):
self.provider = provider
self.config = config or {}
self.logger = structlog.get_logger(self.__class__.__name__)
self.state = {}
@abstractmethod
async def execute(self) -> Dict[str, Any]:
"""Execute the workflow."""
pass
async def run_step(
self,
step: WorkflowStep,
context: Dict = None
) -> str:
"""Execute a single workflow step."""
context = context or {}
self.logger.info(
"Running step",
step=step.name,
task=step.task[:50]
)
# Get LLM response
from ml_agent.providers.base import Message
messages = [
Message(role="user", content=step.task)
]
response = await self.provider.complete(messages)
# Validate if needed
if step.validation:
if not step.validation(response):
raise ValueError(f"Validation failed for step {step.name}")
return response
def _save_state(self, key: str, value: Any):
"""Save state for later steps."""
self.state[key] = value
def _load_state(self, key: str) -> Any:
"""Load saved state."""
return self.state.get(key)- Step 2: Create registry.py
# src/ml_agent/workflows/registry.py
from typing import Dict, Type
from ml_agent.workflows.base import Workflow
class WorkflowRegistry:
"""Registry for available workflows."""
_workflows: Dict[str, Type[Workflow]] = {}
@classmethod
def register(cls, name: str, workflow_class: Type[Workflow]):
"""Register a workflow."""
cls._workflows[name] = workflow_class
@classmethod
def get(cls, name: str) -> Type[Workflow]:
"""Get workflow class."""
if name not in cls._workflows:
raise ValueError(f"Unknown workflow: {name}")
return cls._workflows[name]
@classmethod
def list_available(cls) -> list[str]:
"""List all available workflows."""
return list(cls._workflows.keys())- Step 3: Create init.py
# src/ml_agent/workflows/__init__.py
from ml_agent.workflows.base import Workflow, WorkflowStep
from ml_agent.workflows.registry import WorkflowRegistry
__all__ = ["Workflow", "WorkflowStep", "WorkflowRegistry"]- Step 4: Create test_workflows.py
# tests/unit/test_workflows.py
import pytest
from ml_agent.workflows.base import Workflow, WorkflowStep
from ml_agent.workflows.registry import WorkflowRegistry
class TestWorkflow(Workflow):
name = "test"
description = "Test workflow"
async def execute(self):
return {"status": "success"}
def test_workflow_step():
"""Test workflow step creation."""
step = WorkflowStep(name="test_step", task="Test task")
assert step.name == "test_step"
def test_workflow_registry():
"""Test workflow registry."""
WorkflowRegistry.register("test", TestWorkflow)
assert "test" in WorkflowRegistry.list_available()
assert WorkflowRegistry.get("test") == TestWorkflow
def test_workflow_unknown():
"""Test error on unknown workflow."""
with pytest.raises(ValueError):
WorkflowRegistry.get("unknown")- Step 5: Run tests
pytest tests/unit/test_workflows.py -v
# Expected: PASS- Step 6: Commit
git add src/ml_agent/workflows/ tests/unit/test_workflows.py
git commit -m "feat: add workflow base class and registry"Files:
- Create:
src/ml_agent/workflows/dataset.py
Interfaces:
-
Consumes:
Workflowbase, providers -
Produces:
ArXivDatasetWorkflowclass -
Step 1: Create dataset.py
# src/ml_agent/workflows/dataset.py
import json
from pathlib import Path
from typing import Any, Dict
from ml_agent.workflows.base import Workflow, WorkflowStep
from ml_agent.workflows.registry import WorkflowRegistry
class ArXivDatasetWorkflow(Workflow):
"""Workflow to collect dataset from arXiv papers."""
name = "arxiv-dataset"
description = "Collect LaTeX expressions and explanations from arXiv"
async def execute(self) -> Dict[str, Any]:
"""Execute dataset collection workflow."""
output_file = Path(self.config.get("output_file", "dataset.jsonl"))
self.logger.info("Starting dataset collection", output=str(output_file))
# Step 1: Download papers (simulated)
step1 = WorkflowStep(
name="download_papers",
task="List 10 popular math papers from arXiv in 2024"
)
papers = await self.run_step(step1)
self._save_state("papers", papers)
# Step 2: Extract equations
step2 = WorkflowStep(
name="extract_equations",
task=f"From these papers: {papers[:200]}... extract LaTeX equations"
)
equations = await self.run_step(step2)
self._save_state("equations", equations)
# Step 3: Generate explanations
step3 = WorkflowStep(
name="generate_explanations",
task=f"Generate Portuguese explanations for: {equations[:200]}..."
)
explanations = await self.run_step(step3)
self._save_state("explanations", explanations)
# Step 4: Validate dataset
step4 = WorkflowStep(
name="validate",
task=f"Validate dataset quality: {explanations[:200]}..."
)
validation = await self.run_step(step4)
self.logger.info("Dataset collection complete", validation=validation)
return {
"status": "success",
"papers_processed": len(papers),
"equations_found": len(equations),
"explanations_generated": len(explanations),
"output_file": str(output_file)
}
# Register the workflow
WorkflowRegistry.register("arxiv-dataset", ArXivDatasetWorkflow)- Step 2: Test manually
# Just verify imports work
python -c "from ml_agent.workflows.dataset import ArXivDatasetWorkflow; print('OK')"- Step 3: Commit
git add src/ml_agent/workflows/dataset.py
git commit -m "feat: add ArXiv dataset collection workflow"Files:
- Create:
src/ml_agent/workflows/training.py
Interfaces:
-
Consumes:
Workflowbase, providers -
Produces:
FineTuneWorkflowclass -
Step 1: Create training.py
# src/ml_agent/workflows/training.py
from typing import Any, Dict
from ml_agent.workflows.base import Workflow, WorkflowStep
from ml_agent.workflows.registry import WorkflowRegistry
class FineTuneWorkflow(Workflow):
"""Workflow to fine-tune a model."""
name = "fine-tune"
description = "Fine-tune a model on custom dataset"
async def execute(self) -> Dict[str, Any]:
"""Execute fine-tuning workflow."""
dataset_file = self.config.get("dataset_file", "dataset.jsonl")
model_name = self.config.get("model", "google/flan-t5-small")
self.logger.info("Starting fine-tuning", model=model_name, dataset=dataset_file)
# Step 1: Prepare dataset
step1 = WorkflowStep(
name="prepare_dataset",
task=f"Prepare and validate dataset from {dataset_file}"
)
prep = await self.run_step(step1)
self._save_state("prep", prep)
# Step 2: Setup training
step2 = WorkflowStep(
name="setup_training",
task=f"Setup training environment for {model_name} on prepared dataset"
)
setup = await self.run_step(step2)
self._save_state("setup", setup)
# Step 3: Train model
step3 = WorkflowStep(
name="train",
task=f"Train {model_name} with learning_rate=5e-5, epochs=3"
)
training = await self.run_step(step3)
self._save_state("training", training)
# Step 4: Evaluate
step4 = WorkflowStep(
name="evaluate",
task="Evaluate trained model on test set and report metrics"
)
evaluation = await self.run_step(step4)
self.logger.info("Fine-tuning complete")
return {
"status": "success",
"model": model_name,
"training_metrics": evaluation,
"output_path": "./models/fine-tuned"
}
# Register the workflow
WorkflowRegistry.register("fine-tune", FineTuneWorkflow)- Step 2: Commit
git add src/ml_agent/workflows/training.py
git commit -m "feat: add model fine-tuning workflow"Files:
- Create:
src/ml_agent/workflows/deployment.py
Interfaces:
-
Consumes:
Workflowbase, providers -
Produces:
HubDeploymentWorkflowclass -
Step 1: Create deployment.py
# src/ml_agent/workflows/deployment.py
from typing import Any, Dict
from ml_agent.workflows.base import Workflow, WorkflowStep
from ml_agent.workflows.registry import WorkflowRegistry
class HubDeploymentWorkflow(Workflow):
"""Workflow to deploy model to Hugging Face Hub."""
name = "deploy-hub"
description = "Deploy trained model to Hugging Face Hub"
async def execute(self) -> Dict[str, Any]:
"""Execute deployment workflow."""
model_path = self.config.get("model_path", "./models/fine-tuned")
repo_name = self.config.get("repo_name", "latex-explainer-pt-v1")
self.logger.info("Starting deployment", repo=repo_name, model=model_path)
# Step 1: Prepare model
step1 = WorkflowStep(
name="prepare",
task=f"Prepare {model_path} for publishing to Hub"
)
prep = await self.run_step(step1)
# Step 2: Create repo
step2 = WorkflowStep(
name="create_repo",
task=f"Create Hugging Face Hub repository '{repo_name}' with Apache 2.0 license"
)
repo = await self.run_step(step2)
self._save_state("repo", repo)
# Step 3: Upload model
step3 = WorkflowStep(
name="upload",
task=f"Upload model files to Hub repository {repo_name}"
)
upload = await self.run_step(step3)
# Step 4: Create README
step4 = WorkflowStep(
name="create_readme",
task=f"Create comprehensive README.md with usage examples for {repo_name}"
)
readme = await self.run_step(step4)
self.logger.info("Deployment complete", repo=repo_name)
return {
"status": "success",
"repository": repo_name,
"hub_url": f"https://huggingface.co/your-username/{repo_name}",
"model_id": repo_name
}
# Register the workflow
WorkflowRegistry.register("deploy-hub", HubDeploymentWorkflow)- Step 2: Commit
git add src/ml_agent/workflows/deployment.py
git commit -m "feat: add Hugging Face Hub deployment workflow"Files:
- Create:
src/ml_agent/core/agent.py - Create:
tests/integration/test_agent.py
Interfaces:
-
Consumes: All previous components
-
Produces:
MLAgentmain orchestrator class -
Step 1: Create agent.py
# src/ml_agent/core/agent.py
import asyncio
from typing import Any, Dict, Optional
import structlog
from ml_agent.core.config import Config
from ml_agent.auth.manager import AuthManager
from ml_agent.providers.registry import ProviderRegistry
from ml_agent.workflows.registry import WorkflowRegistry
from ml_agent.core.exceptions import MLAgentException
class MLAgent:
"""Main orchestrator for ML workflows."""
def __init__(
self,
provider: str,
workflow: str,
config: Config = None,
api_key: Optional[str] = None,
workflow_config: Dict[str, Any] = None,
):
self.provider_name = provider
self.workflow_name = workflow
self.config = config or Config()
self.workflow_config = workflow_config or {}
self.logger = structlog.get_logger("MLAgent")
# Setup
self.config.ensure_directories()
# Auth
self.auth_manager = AuthManager(
auth_file=self.config.auth_file,
cli_api_key=api_key
)
# Get provider
provider_key = self.auth_manager.get_api_key(provider)
self.provider = ProviderRegistry.get(provider, provider_key)
# Get workflow
self.workflow_class = WorkflowRegistry.get(workflow)
async def run(self) -> Dict[str, Any]:
"""Run the workflow."""
self.logger.info(
"Starting workflow",
provider=self.provider_name,
workflow=self.workflow_name
)
try:
# Instantiate and execute workflow
workflow = self.workflow_class(
provider=self.provider,
config=self.workflow_config
)
result = await workflow.execute()
self.logger.info("Workflow complete", result=result)
return result
except Exception as e:
self.logger.error("Workflow failed", error=str(e))
raise
def run_sync(self) -> Dict[str, Any]:
"""Run workflow synchronously."""
return asyncio.run(self.run())- Step 2: Create test_agent.py
# tests/integration/test_agent.py
import pytest
from unittest.mock import AsyncMock, patch
from ml_agent.core.agent import MLAgent
from ml_agent.core.config import Config
@pytest.mark.asyncio
async def test_agent_initialization():
"""Test agent initialization."""
with patch('ml_agent.auth.manager.AuthManager.get_api_key', return_value="test-key"):
with patch('ml_agent.providers.registry.ProviderRegistry.get') as mock_get:
mock_provider = AsyncMock()
mock_get.return_value = mock_provider
agent = MLAgent(
provider="claude",
workflow="arxiv-dataset",
)
assert agent.provider_name == "claude"
assert agent.workflow_name == "arxiv-dataset"- Step 3: Run tests
pytest tests/integration/test_agent.py -v
# Expected: PASS- Step 4: Commit
git add src/ml_agent/core/agent.py tests/integration/test_agent.py
git commit -m "feat: add main agent orchestrator"Files:
- Create:
src/ml_agent/plugins/base.py - Create:
src/ml_agent/plugins/loader.py - Create:
tests/unit/test_plugins.py
Interfaces:
-
Produces: Plugin base classes, loader system
-
Step 1: Create base.py
# src/ml_agent/plugins/base.py
from abc import ABC, abstractmethod
from typing import Dict, Any
class Plugin(ABC):
"""Base class for all plugins."""
name: str
version: str
@abstractmethod
def initialize(self, config: Dict[str, Any]) -> None:
"""Initialize plugin with config."""
pass
class ProviderPlugin(Plugin):
"""Plugin for adding a new LLM provider."""
@abstractmethod
def get_provider_class(self):
"""Return provider class."""
pass
class WorkflowPlugin(Plugin):
"""Plugin for adding a new workflow."""
@abstractmethod
def get_workflow_class(self):
"""Return workflow class."""
pass- Step 2: Create loader.py
# src/ml_agent/plugins/loader.py
from pathlib import Path
from typing import List
import importlib.util
import structlog
from ml_agent.plugins.base import Plugin, ProviderPlugin, WorkflowPlugin
from ml_agent.providers.registry import ProviderRegistry
from ml_agent.workflows.registry import WorkflowRegistry
class PluginLoader:
"""Loads and manages plugins."""
def __init__(self, plugin_dir: Path = None):
self.plugin_dir = plugin_dir or Path.home() / ".ml-agent" / "plugins"
self.logger = structlog.get_logger("PluginLoader")
self.loaded_plugins = {}
def load_plugins(self) -> List[Plugin]:
"""Load all plugins from plugin directory."""
if not self.plugin_dir.exists():
return []
plugins = []
for plugin_file in self.plugin_dir.glob("*.py"):
if plugin_file.name.startswith("_"):
continue
try:
plugin = self._load_plugin_file(plugin_file)
if plugin:
plugins.append(plugin)
self.logger.info("Loaded plugin", name=plugin.name)
except Exception as e:
self.logger.error("Failed to load plugin", file=plugin_file, error=str(e))
return plugins
def _load_plugin_file(self, path: Path) -> Plugin:
"""Load single plugin file."""
spec = importlib.util.spec_from_file_location(path.stem, path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# Find Plugin class
for attr_name in dir(module):
attr = getattr(module, attr_name)
if isinstance(attr, type) and issubclass(attr, Plugin) and attr != Plugin:
plugin = attr()
# Register appropriately
if isinstance(plugin, ProviderPlugin):
provider_class = plugin.get_provider_class()
ProviderRegistry.register(plugin.name, provider_class)
elif isinstance(plugin, WorkflowPlugin):
workflow_class = plugin.get_workflow_class()
WorkflowRegistry.register(plugin.name, workflow_class)
return plugin
return None- Step 3: Create test_plugins.py
# tests/unit/test_plugins.py
from ml_agent.plugins.base import Plugin, ProviderPlugin, WorkflowPlugin
class MockPlugin(Plugin):
name = "mock"
version = "1.0.0"
def initialize(self, config):
pass
def test_plugin_base():
"""Test plugin base class."""
plugin = MockPlugin()
assert plugin.name == "mock"
assert plugin.version == "1.0.0"- Step 4: Commit
git add src/ml_agent/plugins/ tests/unit/test_plugins.py
git commit -m "feat: add plugin system for extensibility"Files:
- Create:
tests/integration/test_end_to_end.py - Create:
docs/API_REFERENCE.md - Create:
docs/PLUGINS.md - Modify:
README.md
Interfaces:
-
Consumes: All components
-
Produces: Comprehensive tests and docs
-
Step 1: Create end-to-end test
# tests/integration/test_end_to_end.py
import pytest
from unittest.mock import AsyncMock, patch
from ml_agent.core.agent import MLAgent
@pytest.mark.asyncio
@patch('ml_agent.providers.registry.ProviderRegistry.get')
@patch('ml_agent.auth.manager.AuthManager.get_api_key')
async def test_full_workflow(mock_auth, mock_provider):
"""Test complete workflow execution."""
# Setup mocks
mock_auth.return_value = "test-key"
mock_provider_instance = AsyncMock()
mock_provider_instance.complete = AsyncMock(return_value="Mocked response")
mock_provider.return_value = mock_provider_instance
# Run agent
agent = MLAgent(
provider="claude",
workflow="arxiv-dataset",
workflow_config={"max_papers": 5}
)
# This would execute the full workflow
# result = agent.run_sync()
# assert result["status"] == "success"- Step 2: Create API_REFERENCE.md
# ML Agent Framework - API Reference
## Core Classes
### MLAgent
Main orchestrator for ML workflows.
```python
from ml_agent.core.agent import MLAgent
agent = MLAgent(
provider="claude",
workflow="arxiv-dataset",
workflow_config={"output_file": "dataset.jsonl"}
)
result = agent.run_sync()Base class for workflows.
from ml_agent.workflows.base import Workflow
class MyWorkflow(Workflow):
name = "my-workflow"
async def execute(self):
step = WorkflowStep(...)
await self.run_step(step)claude- Anthropic Claudeopenai- OpenAI GPTdeepseek- DeepSeekmistral- Mistral AI
arxiv-dataset- Collect dataset from arXivfine-tune- Fine-tune modeldeploy-hub- Deploy to Hugging Face Hub
- [ ] **Step 3: Create PLUGINS.md**
```markdown
# Plugin System
## Creating a Provider Plugin
```python
from ml_agent.plugins.base import ProviderPlugin
from ml_agent.providers.base import BaseProvider
class MyProvider(BaseProvider):
name = "myprovider"
async def complete(self, messages, **kwargs):
# Implementation
pass
class MyProviderPlugin(ProviderPlugin):
name = "myprovider-plugin"
version = "1.0.0"
def initialize(self, config):
pass
def get_provider_class(self):
return MyProvider
Place in ~/.ml-agent/plugins/my_provider.py
- [ ] **Step 4: Update README.md**
```markdown
# ML Agent Framework
Multi-provider ML workflow automation framework supporting Claude, OpenAI, DeepSeek, and Mistral.
## Installation
```bash
pip install ml-agent
# List providers
ml-agent list-providers
# Validate credentials
ml-agent validate-auth --provider claude
# Run workflow
ml-agent run \
--provider claude \
--workflow arxiv-dataset \
--config config.yamlApache 2.0
- [ ] **Step 5: Run all tests**
```bash
pytest tests/ -v --cov=src/ml_agent
# Expected: >80% coverage
- Step 6: Final commit
git add tests/integration/ docs/ README.md
git commit -m "feat: add integration tests and documentation"All 16 tasks complete!
✅ Phase 1-2: Foundation & Core (7 tasks)
- Project setup, config, logging
- 2 LLM providers (Claude, OpenAI)
- Provider registry
✅ Phase 3: Auth & CLI (2 tasks)
- Multi-strategy authentication
- Typer CLI with commands
✅ Phase 4: Workflows (4 tasks)
- Workflow base & registry
- Dataset collection (arXiv)
- Model fine-tuning
- Hub deployment
✅ Phase 5: Integration (3 tasks)
- Agent orchestrator
- Plugin system
- E2E tests & docs
- ~2,000 lines of production code
- ~1,500 lines of test code
- Comprehensive documentation
- Ready for open-source release
- Add DeepSeek & Mistral providers
- Create example workflows
- Set up CI/CD in GitHub Actions
- Release on PyPI
- Launch in p7dotorg
Ready to start implementation! 🚀