code-muse 0.0.1__py3-none-any.whl

code_muse/plugins/azure_foundry/config.py

@@ -0,0 +1,125 @@
+"""Configuration constants for the Azure AI Foundry plugin.
+
+This module defines constants used throughout the plugin for Azure AD
+authentication and model configuration.
+"""
+
+import os
+from pathlib import Path
+from typing import Any
+
+from code_muse.config import DATA_DIR
+
+# Azure AD scope for Cognitive Services (used for token acquisition)
+# This scope is required for authenticating with Azure AI Foundry
+AZURE_COGNITIVE_SCOPE = "https://cognitiveservices.azure.com/.default"
+
+# Token refresh buffer in seconds (refresh 5 minutes before expiry)
+TOKEN_REFRESH_BUFFER = 300
+
+# Default context lengths for different model tiers
+DEFAULT_CONTEXT_LENGTHS: dict[str, int] = {
+    "opus": 1000000,  # 1M tokens for Opus models
+    "sonnet": 1000000,  # 1M tokens for Sonnet models
+    "haiku": 200000,  # 200K tokens for Haiku models
+}
+
+# Context lengths for OpenAI models (Azure doesn't expose this in the catalog API).
+# Prefixes are matched longest-first against the model name.
+OPENAI_CONTEXT_LENGTHS: dict[str, int] = {
+    "gpt-5.4": 1000000,
+    "gpt-5.4-mini": 1000000,
+    "gpt-5.3-codex": 1000000,
+    "gpt-5.3": 1000000,
+    "gpt-5.2-codex": 1000000,
+    "gpt-5.2": 1000000,
+    "gpt-5.1-codex-max": 1000000,
+    "gpt-5.1-codex-mini": 1000000,
+    "gpt-5.1-codex": 1000000,
+    "gpt-5.1": 1000000,
+    "gpt-5-codex": 1000000,
+    "gpt-5": 1000000,
+    "gpt-4.1": 1000000,
+    "gpt-4.1-mini": 1000000,
+    "gpt-4.1-nano": 1000000,
+    "o4-mini": 200000,
+    "o3": 200000,
+    "o3-mini": 200000,
+    "o1": 200000,
+    "o1-mini": 128000,
+    "codex-mini": 200000,
+}
+DEFAULT_OPENAI_CONTEXT_LENGTH = 128000
+
+
+def get_openai_context_length(model_name: str) -> int:
+    """Look up the context length for an OpenAI model by name.
+
+    Matches the longest prefix first so 'gpt-5.4-mini' matches before 'gpt-5.4'.
+    Falls back to DEFAULT_OPENAI_CONTEXT_LENGTH if no match.
+    """
+    for prefix in sorted(OPENAI_CONTEXT_LENGTHS, key=len, reverse=True):
+        if model_name.startswith(prefix):
+            return OPENAI_CONTEXT_LENGTHS[prefix]
+    return DEFAULT_OPENAI_CONTEXT_LENGTH
+
+
+# Default deployment name patterns (can be overridden by user)
+DEFAULT_DEPLOYMENT_NAMES: dict[str, str] = {
+    "opus": "claude-opus-4-6",
+    "sonnet": "claude-sonnet-4-6",
+    "haiku": "claude-haiku-4-5",
+}
+
+# Environment variable names
+ENV_FOUNDRY_RESOURCE = "ANTHROPIC_FOUNDRY_RESOURCE"
+ENV_FOUNDRY_BASE_URL = "ANTHROPIC_FOUNDRY_BASE_URL"
+
+
+def get_foundry_resource() -> str | None:
+    """Get the Azure Foundry resource name from environment.
+
+    Returns:
+        The resource name if set, None otherwise.
+    """
+    return os.environ.get(ENV_FOUNDRY_RESOURCE)
+
+
+def get_foundry_base_url() -> str | None:
+    """Get the Azure Foundry base URL from environment.
+
+    If ANTHROPIC_FOUNDRY_BASE_URL is set, use it directly.
+    Otherwise, construct from ANTHROPIC_FOUNDRY_RESOURCE.
+
+    Returns:
+        The base URL if determinable, None otherwise.
+    """
+    base_url = os.environ.get(ENV_FOUNDRY_BASE_URL)
+    if base_url:
+        return base_url
+
+    resource = get_foundry_resource()
+    if resource:
+        return f"https://{resource}.services.ai.azure.com/anthropic/v1"
+
+    return None
+
+
+def get_extra_models_path() -> Path:
+    """Get the path to the extra_models.json file.
+
+    Returns:
+        Path to the extra models configuration file.
+    """
+    return Path(DATA_DIR) / "extra_models.json"
+
+
+# Plugin configuration for model creation
+AZURE_FOUNDRY_CONFIG: dict[str, Any] = {
+    "scope": AZURE_COGNITIVE_SCOPE,
+    "token_refresh_buffer": TOKEN_REFRESH_BUFFER,
+}
+
+# Note: supported_settings is determined per-tier in build_foundry_model_config():
+# - Opus 4.6: ["temperature", "extended_thinking", "effort"]
+# - Sonnet/Haiku: ["temperature", "interleaved_thinking", "thinking_budget"]

code_muse/plugins/azure_foundry/discovery.py

@@ -0,0 +1,187 @@
+"""Azure AI Services deployment discovery.
+
+Queries the Azure Management API to find AI Services accounts and
+list their model deployments. Works with any AIServices account
+hosting Anthropic, OpenAI, or other model formats.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+AZURE_MANAGEMENT_SCOPE = "https://management.azure.com/.default"
+MANAGEMENT_BASE = "https://management.azure.com"
+RESOURCE_API_VERSION = "2021-04-01"
+DEPLOYMENT_API_VERSION = "2024-10-01"
+
+
+@dataclass
+class AzureAccount:
+    """Discovered Azure AI Services account."""
+
+    resource_id: str
+    name: str
+    location: str
+    resource_group: str
+    subscription_id: str
+
+
+@dataclass
+class AzureDeployment:
+    """Discovered model deployment on an Azure AI Services account."""
+
+    name: str
+    model_name: str
+    model_format: str  # "Anthropic", "OpenAI", etc.
+    model_version: str
+    provisioning_state: str
+    sku_name: str
+    capacity: int
+
+
+def _get_management_token() -> str | None:
+    """Get a token for the Azure Management API using az login credentials.
+
+    Returns:
+        Bearer token string, or None if auth fails.
+    """
+    try:
+        from azure.identity import AzureCliCredential
+
+        credential = AzureCliCredential()
+        token = credential.get_token(AZURE_MANAGEMENT_SCOPE)
+        return token.token
+    except Exception as e:
+        logger.warning("Failed to get management token: %s", e)
+        return None
+
+
+def _management_get(token: str, url: str) -> dict[str, Any] | None:
+    """Make a GET request to the Azure Management API.
+
+    Args:
+        token: Bearer token for authentication.
+        url: Full URL to GET.
+
+    Returns:
+        Parsed JSON response, or None on error.
+    """
+    try:
+        import httpx
+
+        resp = httpx.get(
+            url,
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=30,
+        )
+        if resp.status_code == 200:
+            return resp.json()
+        logger.warning("Management API %s returned %d", url, resp.status_code)
+        return None
+    except Exception as e:
+        logger.warning("Management API request failed: %s", e)
+        return None
+
+
+def find_account(resource_name: str) -> AzureAccount | None:
+    """Find an Azure AI Services account by name across all accessible subscriptions.
+
+    Args:
+        resource_name: The account name to search for.
+
+    Returns:
+        AzureAccount if found, None otherwise.
+    """
+    token = _get_management_token()
+    if not token:
+        return None
+
+    # List subscriptions
+    subs_url = f"{MANAGEMENT_BASE}/subscriptions?api-version=2022-12-01"
+    subs_resp = _management_get(token, subs_url)
+    if not subs_resp:
+        return None
+
+    subscriptions = subs_resp.get("value", [])
+
+    for sub in subscriptions:
+        sub_id = sub.get("subscriptionId", "")
+        if sub.get("state") != "Enabled":
+            continue
+
+        # Search for the resource by name and type
+        filter_str = (
+            f"name eq '{resource_name}' and "
+            f"resourceType eq 'Microsoft.CognitiveServices/accounts'"
+        )
+        resources_url = (
+            f"{MANAGEMENT_BASE}/subscriptions/{sub_id}/resources"
+            f"?$filter={filter_str}&api-version={RESOURCE_API_VERSION}"
+        )
+        resources_resp = _management_get(token, resources_url)
+        if not resources_resp:
+            continue
+
+        for resource in resources_resp.get("value", []):
+            rid = resource.get("id", "")
+            parts = rid.split("/")
+            # /subscriptions/{sub}/resourceGroups/{rg}/providers/.../accounts/{name}
+            rg_idx = next(
+                (i for i, p in enumerate(parts) if p.lower() == "resourcegroups"),
+                None,
+            )
+            rg = parts[rg_idx + 1] if rg_idx is not None else ""
+
+            return AzureAccount(
+                resource_id=rid,
+                name=resource_name,
+                location=resource.get("location", ""),
+                resource_group=rg,
+                subscription_id=sub_id,
+            )
+
+    return None
+
+
+def list_deployments(account: AzureAccount) -> list[AzureDeployment]:
+    """List all model deployments on an Azure AI Services account.
+
+    Args:
+        account: The account to query.
+
+    Returns:
+        List of deployments (all states, caller filters as needed).
+    """
+    token = _get_management_token()
+    if not token:
+        return []
+
+    url = (
+        f"{MANAGEMENT_BASE}{account.resource_id}"
+        f"/deployments?api-version={DEPLOYMENT_API_VERSION}"
+    )
+    resp = _management_get(token, url)
+    if not resp:
+        return []
+
+    deployments = []
+    for d in resp.get("value", []):
+        props = d.get("properties", {})
+        model = props.get("model", {})
+        sku = d.get("sku", {})
+
+        deployments.append(
+            AzureDeployment(
+                name=d.get("name", ""),
+                model_name=model.get("name", ""),
+                model_format=model.get("format", ""),
+                model_version=model.get("version", ""),
+                provisioning_state=props.get("provisioningState", ""),
+                sku_name=sku.get("name", ""),
+                capacity=sku.get("capacity", 0),
+            )
+        )
+
+    return deployments