code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/aws_bedrock/register_callbacks.py

@@ -0,0 +1,241 @@
+"""AWS Bedrock Plugin callbacks for Muse CLI.
+
+This plugin enables Muse to use Anthropic Claude models hosted on
+AWS Bedrock with standard AWS credential chain authentication.
+"""
+
+import logging
+from typing import Any
+
+from code_muse.callbacks import register_callback
+from code_muse.messaging import emit_error, emit_info, emit_success, emit_warning
+
+from .config import (
+    get_aws_profile,
+    get_bedrock_region,
+)
+from .utils import (
+    add_bedrock_models_to_config,
+    get_bedrock_models_from_config,
+    remove_bedrock_models_from_config,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# ============================================================================
+# Slash Command Handlers
+# ============================================================================
+
+
+def _handle_bedrock_status() -> None:
+    """Handle the /bedrock-status command."""
+    emit_info("")
+    emit_info("AWS Bedrock Status")
+    emit_info("=" * 40)
+
+    region = get_bedrock_region()
+    profile = get_aws_profile()
+
+    emit_info(f"Region: {region}")
+    if profile:
+        emit_info(f"Profile: {profile}")
+    else:
+        emit_info("Profile: (default credential chain)")
+
+    # Check AWS credentials
+    try:
+        import boto3
+
+        session = boto3.Session(profile_name=profile, region_name=region)
+        sts = session.client("sts")
+        identity = sts.get_caller_identity()
+        arn = identity.get("Arn", "unknown")
+        emit_success(f"Authenticated: {arn}")
+    except ImportError:
+        emit_warning("boto3 not installed - cannot verify credentials")
+    except Exception as e:
+        emit_warning(f"Authentication: {e}")
+
+    emit_info("")
+    bedrock_models = get_bedrock_models_from_config()
+    if bedrock_models:
+        emit_info(f"Configured Models ({len(bedrock_models)}):")
+        for model_key, config in bedrock_models.items():
+            model_id = config.get("name", "unknown")
+            emit_info(f"   - {model_key}: {model_id}")
+    else:
+        emit_info("Configured Models: None")
+        emit_info("   Run /bedrock-setup to configure models")
+
+    emit_info("")
+
+
+def _handle_bedrock_setup() -> None:
+    """Handle the /bedrock-setup command.
+
+    Region and credentials auto-resolve from the AWS environment
+    (IAM role, env vars, ~/.aws/config). No prompts needed.
+    """
+
+    added_models = add_bedrock_models_to_config(
+        aws_region=get_bedrock_region(),
+        aws_profile=get_aws_profile(),
+    )
+
+    if added_models:
+        emit_success(f"Bedrock: added {len(added_models)} model(s):")
+        for model_key in added_models:
+            emit_info(f"   - {model_key}")
+
+        from code_muse.model_switching import set_model_and_reload_agent
+
+        set_model_and_reload_agent("bedrock-opus-4-7")
+        emit_info("Switched to bedrock-opus-4-7.")
+    else:
+        emit_warning("Bedrock: no models added — check configuration.")
+
+
+def _handle_bedrock_remove() -> None:
+    """Handle the /bedrock-remove command."""
+    removed = remove_bedrock_models_from_config()
+    if removed:
+        emit_success(f"Removed {len(removed)} Bedrock model(s):")
+        for model_key in removed:
+            emit_info(f"   - {model_key}")
+    else:
+        emit_info("No Bedrock models found in configuration.")
+
+
+# ============================================================================
+# Custom Command Registration
+# ============================================================================
+
+
+def _custom_help() -> list[tuple[str, str]]:
+    """Return help entries for custom commands."""
+    return [
+        ("bedrock-status", "Check AWS Bedrock authentication and configuration"),
+        ("bedrock-setup", "Auto-configure Bedrock Claude models"),
+        ("bedrock-remove", "Remove all Bedrock model configurations"),
+    ]
+
+
+def _handle_custom_command(command: str, name: str) -> bool | None:
+    """Handle custom slash commands for the AWS Bedrock plugin."""
+    handlers = {
+        "bedrock-status": _handle_bedrock_status,
+        "bedrock-setup": _handle_bedrock_setup,
+        "bedrock-remove": _handle_bedrock_remove,
+    }
+
+    handler = handlers.get(name)
+    if handler is None:
+        return None
+
+    try:
+        handler()
+        return True
+    except Exception as e:
+        logger.exception("Error handling /%s command: %s", name, e)
+        emit_error(f"Command /{name} failed: {e}")
+        return True
+
+
+# ============================================================================
+# Model Type Handler
+# ============================================================================
+
+
+def _create_aws_bedrock_model(model_name: str, model_config: dict, config: dict) -> Any:
+    """Create an AWS Bedrock model instance.
+
+    Uses AsyncAnthropicBedrock from the anthropic SDK with standard
+    AWS credential chain (env vars, profiles, IAM roles, SSO).
+    """
+    try:
+        from anthropic import AsyncAnthropicBedrock
+        from pydantic_ai.models.anthropic import AnthropicModel
+    except ImportError as e:
+        emit_error(
+            f"Failed to create Bedrock model '{model_name}': Missing dependency - {e}"
+        )
+        return None
+
+    from code_muse.claude_cache_client import patch_anthropic_client_messages
+    from code_muse.config import get_effective_model_settings
+    from code_muse.model_factory import CONTEXT_1M_BETA
+    from code_muse.provider_identity import (
+        make_anthropic_provider,
+        resolve_provider_identity,
+    )
+
+    model_id = model_config.get("name")
+    if not model_id:
+        emit_warning(f"Model ID not specified for '{model_name}'.")
+        return None
+
+    try:
+        effective_settings = get_effective_model_settings(model_name)
+        interleaved_thinking = effective_settings.get("interleaved_thinking", True)
+
+        beta_parts: list[str] = []
+        if interleaved_thinking:
+            beta_parts.append("interleaved-thinking-2025-05-14")
+
+        context_length = model_config.get("context_length", 200000)
+        if context_length >= 1_000_000:
+            beta_parts.append(CONTEXT_1M_BETA)
+
+        default_headers: dict[str, str] | None = None
+        if beta_parts:
+            default_headers = {"anthropic-beta": ",".join(beta_parts)}
+
+        client_kwargs: dict[str, Any] = {
+            "aws_region": model_config.get("aws_region") or get_bedrock_region(),
+        }
+
+        aws_profile = model_config.get("aws_profile") or get_aws_profile()
+        if aws_profile:
+            client_kwargs["aws_profile"] = aws_profile
+
+        if default_headers:
+            client_kwargs["default_headers"] = default_headers
+
+        anthropic_client = AsyncAnthropicBedrock(**client_kwargs)
+
+        patch_anthropic_client_messages(anthropic_client)
+
+        provider_identity = resolve_provider_identity(model_name, model_config)
+        provider = make_anthropic_provider(
+            provider_identity,
+            anthropic_client=anthropic_client,
+        )
+
+        model = AnthropicModel(model_name=model_id, provider=provider)
+        logger.info(
+            "Created Bedrock model: %s -> %s @ %s",
+            model_name,
+            model_id,
+            client_kwargs.get("aws_region", "auto"),
+        )
+        return model
+
+    except Exception as e:
+        emit_error(f"Failed to create Bedrock model '{model_name}': {e}")
+        logger.exception("Error creating Bedrock model: %s", e)
+        return None
+
+
+def _register_model_types() -> list[dict[str, Any]]:
+    """Register the aws_bedrock model type handler."""
+    return [{"type": "aws_bedrock", "handler": _create_aws_bedrock_model}]
+
+
+# ============================================================================
+# Callback Registration
+# ============================================================================
+
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)
+register_callback("register_model_type", _register_model_types)

code_muse/plugins/aws_bedrock/utils.py

@@ -0,0 +1,153 @@
+"""Utility functions for the AWS Bedrock plugin."""
+
+import json
+import logging
+from typing import Any
+
+from .config import (
+    MODELS,
+    get_extra_models_path,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def load_extra_models() -> dict[str, Any]:
+    """Load the extra_models.json configuration file."""
+    extra_models_path = get_extra_models_path()
+    if not extra_models_path.exists():
+        return {}
+
+    try:
+        with open(extra_models_path, encoding="utf-8") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError) as e:
+        logger.error("Error loading extra_models.json: %s", e)
+        return {}
+
+
+def save_extra_models(models: dict[str, Any]) -> bool:
+    """Save model configurations to extra_models.json (atomic write)."""
+    extra_models_path = get_extra_models_path()
+
+    try:
+        extra_models_path.parent.mkdir(parents=True, exist_ok=True)
+        temp_path = extra_models_path.with_suffix(".tmp")
+        with open(temp_path, "w", encoding="utf-8") as f:
+            json.dump(models, f, indent=2, ensure_ascii=False)
+        temp_path.replace(extra_models_path)
+        return True
+    except Exception as e:
+        logger.error("Error saving extra_models.json: %s", e)
+        return False
+
+
+def _build_model_entry(
+    model_id: str,
+    context_length: int,
+    has_effort: bool,
+    effort: str | None = None,
+    aws_region: str | None = None,
+    aws_profile: str | None = None,
+) -> dict[str, Any]:
+    """Build a single model config entry for extra_models.json."""
+    supported_settings = [
+        "temperature",
+        "extended_thinking",
+        "budget_tokens",
+        "interleaved_thinking",
+    ]
+    if has_effort:
+        supported_settings.append("effort")
+
+    config: dict[str, Any] = {
+        "type": "aws_bedrock",
+        "provider": "aws_bedrock",
+        "name": model_id,
+        "context_length": context_length,
+        "supported_settings": supported_settings,
+    }
+    if effort:
+        config["default_effort"] = effort
+    if aws_region:
+        config["aws_region"] = aws_region
+    if aws_profile:
+        config["aws_profile"] = aws_profile
+
+    return config
+
+
+def add_bedrock_models_to_config(
+    aws_region: str | None = None,
+    aws_profile: str | None = None,
+) -> list[str]:
+    """Add Bedrock model configurations (with effort variants) to extra_models.json."""
+    models = load_extra_models()
+    added: list[str] = []
+
+    for spec in MODELS:
+        base_key = spec["base_key"]
+        model_id = spec["model_id"]
+        context_length = spec["context_length"]
+        variants = spec.get("variants")
+
+        if variants:
+            for variant in variants:
+                if variant == "default":
+                    key = base_key
+                    effort = None
+                else:
+                    key = f"{base_key}-{variant}"
+                    effort = variant
+
+                models[key] = _build_model_entry(
+                    model_id=model_id,
+                    context_length=context_length,
+                    has_effort=True,
+                    effort=effort,
+                    aws_region=aws_region,
+                    aws_profile=aws_profile,
+                )
+                added.append(key)
+        else:
+            models[base_key] = _build_model_entry(
+                model_id=model_id,
+                context_length=context_length,
+                has_effort=False,
+                aws_region=aws_region,
+                aws_profile=aws_profile,
+            )
+            added.append(base_key)
+
+    if added and save_extra_models(models):
+        return added
+    return []
+
+
+def remove_bedrock_models_from_config() -> list[str]:
+    """Remove all Bedrock model configurations from extra_models.json."""
+    models = load_extra_models()
+    removed = [
+        key
+        for key, cfg in models.items()
+        if isinstance(cfg, dict) and cfg.get("type") == "aws_bedrock"
+    ]
+
+    for key in removed:
+        del models[key]
+
+    if removed and not save_extra_models(models):
+        logger.error("Failed to save extra_models.json after removing Bedrock models")
+        return []
+
+    return removed
+
+
+def get_bedrock_models_from_config() -> dict[str, Any]:
+    """Get all Bedrock model configurations from extra_models.json."""
+    models = load_extra_models()
+    return {
+        key: cfg
+        for key, cfg in models.items()
+        if isinstance(cfg, dict) and cfg.get("type") == "aws_bedrock"
+    }

code_muse/plugins/azure_foundry/README.md

@@ -0,0 +1,238 @@
+# Azure AI Foundry Plugin for Muse
+
+This plugin enables Muse to use Anthropic Claude models hosted on
+Microsoft Azure AI Foundry with Azure AD (Entra ID) authentication.
+
+## Overview
+
+Azure AI Foundry provides enterprise-grade hosting of Anthropic Claude models
+within the Microsoft Azure cloud. This plugin uses your existing Azure CLI
+credentials (`az login`) to authenticate, eliminating the need for API keys.
+
+## Supported Models
+
+| Model | Context Length | Deployment Name (Default) |
+|-------|----------------|---------------------------|
+| Claude Opus 4.6 | 1M tokens | `claude-opus-4-6` |
+| Claude Sonnet 4.6 | 1M tokens | `claude-sonnet-4-6` |
+| Claude Haiku 4.5 | 200K tokens | `claude-haiku-4-5` |
+
+Deployment names are user-configurable during setup.
+
+## Prerequisites
+
+1. **Azure subscription** with access to Azure AI Foundry
+2. **Azure CLI** installed and authenticated (`az login`)
+3. **Claude model deployments** provisioned in your Azure AI Foundry resource
+4. **Python packages**: `azure-identity>=1.15.0` (installed with Muse)
+
+## Quick Start
+
+### 1. Authenticate with Azure
+
+```bash
+az login
+```
+
+### 2. Set Your Resource Name
+
+```bash
+export ANTHROPIC_FOUNDRY_RESOURCE=your-resource-name
+```
+
+### 3. Run Interactive Setup
+
+```bash
+pup
+/foundry-setup
+```
+
+The wizard will guide you through configuring your model deployments.
+
+### 4. Start Using Foundry Models
+
+```bash
+/model foundry-claude-opus
+```
+
+## Configuration
+
+### Environment Variables
+
+| Variable | Description | Required |
+|----------|-------------|----------|
+| `ANTHROPIC_FOUNDRY_RESOURCE` | Your Azure AI Foundry resource name | Yes |
+| `ANTHROPIC_FOUNDRY_BASE_URL` | Override the base URL (optional) | No |
+
+### Model Configuration
+
+Models are stored in `~/.muse/extra_models.json`:
+
+```json
+{
+  "foundry-claude-opus": {
+    "type": "azure_foundry",
+    "name": "it-entra-claude-opus-4-6[1m]",
+    "foundry_resource": "$ANTHROPIC_FOUNDRY_RESOURCE",
+    "context_length": 1000000,
+    "supported_settings": ["temperature", "extended_thinking", "effort"]
+  },
+  "foundry-claude-sonnet": {
+    "type": "azure_foundry",
+    "name": "it-entra-claude-sonnet-4-6[1m]",
+    "foundry_resource": "$ANTHROPIC_FOUNDRY_RESOURCE",
+    "context_length": 1000000,
+    "supported_settings": ["temperature", "interleaved_thinking", "thinking_budget"]
+  },
+  "foundry-claude-haiku": {
+    "type": "azure_foundry",
+    "name": "it-entra-claude-haiku-4-5",
+    "foundry_resource": "$ANTHROPIC_FOUNDRY_RESOURCE",
+    "context_length": 200000,
+    "supported_settings": ["temperature", "interleaved_thinking", "thinking_budget"]
+  }
+}
+```
+
+**Configuration Fields:**
+
+- `type`: Must be `"azure_foundry"` to route to this plugin
+- `name`: Your Azure deployment name (e.g., `it-entra-claude-opus-4-6`)
+- `foundry_resource`: Resource name (supports `$ENV_VAR` syntax)
+- `context_length`: Model context window in tokens
+- `supported_settings`: List of supported model settings
+
+## Model Settings by Tier
+
+### Opus 4.6 Models
+- `temperature`: Temperature for sampling (0-1)
+- `extended_thinking`: Enable thinking (defaults to "adaptive")
+- `effort`: Thinking effort level ("low", "medium", "high")
+
+### Sonnet & Haiku Models
+- `temperature`: Temperature for sampling (0-1)
+- `interleaved_thinking`: Enable interleaved thinking (boolean)
+- `thinking_budget`: Maximum tokens for internal reasoning
+
+## Context Window Configuration
+
+Deployment names can include a context window suffix following the Claude Code format:
+- `[1m]` - 1 million tokens
+- `[2m]` - 2 million tokens
+- `[200k]` - 200 thousand tokens
+- `[500k]` - 500 thousand tokens
+
+When you specify a deployment name with a context suffix (e.g., `it-entra-claude-opus-4-6[1m]`):
+1. The suffix is **automatically stripped** before sending to Azure
+2. The `context_length` is set based on the parsed suffix
+
+**Example:**
+
+During `/foundry-setup`, if you enter:
+```
+Opus deployment name: it-entra-claude-opus-4-6[1m]
+```
+
+The saved configuration will have:
+- `name`: `it-entra-claude-opus-4-6` (suffix stripped)
+- `context_length`: `1000000` (parsed from `[1m]`)
+
+This is useful when your Azure deployment names don't include the context indicator,
+but you want to configure the context length.
+
+## Slash Commands
+
+| Command | Description |
+|---------|-------------|
+| `/foundry-status` | Check Azure AD authentication status and configured models |
+| `/foundry-setup` | Interactive wizard to configure Foundry models |
+| `/foundry-remove` | Remove all Foundry model configurations |
+
+### Example: /foundry-status
+
+```
+Azure AI Foundry Status
+========================================
+Authentication: Valid (expires in 45 minutes)
+   Logged in as: user@company.com
+
+Foundry Resource: my-foundry-resource
+
+Configured Models (3):
+   - foundry-claude-opus: it-entra-claude-opus-4-6[1m]
+   - foundry-claude-sonnet: it-entra-claude-sonnet-4-6[1m]
+   - foundry-claude-haiku: it-entra-claude-haiku-4-5
+```
+
+## How It Works
+
+1. **Token Acquisition**: Uses `AzureCliCredential` from `azure-identity` to
+   obtain tokens from your `az login` session
+
+2. **Token Refresh**: The `get_bearer_token_provider` function handles automatic
+   token refresh before expiry
+
+3. **API Calls**: Creates an `AsyncAnthropicFoundry` client that uses the native
+   Anthropic Messages API (not OpenAI format)
+
+4. **Integration**: Wraps the client in pydantic-ai's `AnthropicModel` for
+   seamless integration with Muse's agent system
+
+## Troubleshooting
+
+### "CredentialUnavailableError" or "Not authenticated"
+
+Run `az login` to authenticate with Azure:
+
+```bash
+az login
+```
+
+### "Resource not found" or "404 errors"
+
+1. Verify `ANTHROPIC_FOUNDRY_RESOURCE` is set correctly
+2. Check that your model deployments exist in Azure AI Foundry
+3. Ensure the deployment names match your configuration
+
+### Token Expiry
+
+Tokens are automatically refreshed by the Azure Identity library. If you
+encounter issues, try:
+
+```bash
+az account get-access-token --resource https://cognitiveservices.azure.com
+```
+
+### Model Not Found
+
+After running `/foundry-setup`, verify configuration with `/foundry-status`.
+Check that deployment names match exactly (case-sensitive).
+
+## Architecture
+
+```
+code_muse/plugins/azure_foundry/
+├── __init__.py              # Package marker and version
+├── config.py                # Constants and configuration helpers
+├── token.py                 # Azure AD token provider
+├── utils.py                 # Configuration file utilities
+└── register_callbacks.py    # Plugin callbacks and model handler
+```
+
+## Security Notes
+
+- **No API keys stored**: Authentication uses Azure AD tokens only
+- **Token caching**: Managed by Azure CLI, not stored by this plugin
+- **Credential scope**: Limited to `https://cognitiveservices.azure.com/.default`
+
+## Contributing
+
+This plugin follows Muse's plugin architecture. Key callbacks:
+
+- `register_model_type`: Registers `azure_foundry` type handler
+- `custom_command`: Handles `/foundry-*` slash commands
+- `custom_command_help`: Provides help text for commands
+
+## License
+
+Same as Muse (see repository LICENSE file).

code_muse/plugins/azure_foundry/__init__.py

@@ -0,0 +1,15 @@
+"""Azure AI Foundry Plugin for Muse.
+
+This plugin enables Muse to use Anthropic Claude models hosted on
+Microsoft Azure AI Foundry with Azure AD (Entra ID) authentication.
+
+The plugin uses the `az login` credentials from the Azure CLI to authenticate,
+eliminating the need for API keys.
+
+Supported models:
+- Claude Opus 4.6 (1M context)
+- Claude Sonnet 4.6 (1M context)
+- Claude Haiku 4.5 (200K context)
+"""
+
+__version__ = "0.1.0"