beaver_project/app-instance/backend/beaver/foundation/config/schema.py

"""Runtime configuration schema for Beaver sandbox instances."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass(slots=True)
class ProviderConfig:
    """One configured LLM provider profile."""

    api_key: str | None = None
    api_base: str | None = None
    extra_headers: dict[str, str] = field(default_factory=dict)
    request_timeout_seconds: float | None = None


@dataclass(slots=True)
class AgentDefaultsConfig:
    """Default agent settings for this sandbox instance."""

    workspace: str | None = None
    model: str | None = None
    provider: str | None = None
    embedding_model: str | None = None
    max_tokens: int | None = None
    temperature: float | None = None
    max_context_messages: int | None = None
    max_tool_iterations: int | None = None


@dataclass(slots=True)
class EmbeddingConfig:
    """Optional dedicated embedding model settings."""

    provider: str | None = None
    model: str | None = None
    api_key: str | None = None
    api_base: str | None = None
    extra_headers: dict[str, str] = field(default_factory=dict)
    request_timeout_seconds: float | None = None


@dataclass(slots=True)
class MCPServerConfig:
    """One configured MCP server.

    Transport is inferred from fields:
    - command => local stdio MCP server
    - url => remote streamable HTTP MCP server
    """

    command: str = ""
    args: list[str] = field(default_factory=list)
    env: dict[str, str] = field(default_factory=dict)
    url: str = ""
    headers: dict[str, str] = field(default_factory=dict)
    auth_mode: str = "none"
    auth_audience: str = ""
    auth_scopes: list[str] = field(default_factory=list)
    tool_timeout: int = 30
    sensitive: bool = False
    kind: str = "online"
    category: str = "online"
    managed: bool = False
    display_name: str = ""
    source: str = "config"

    @property
    def transport(self) -> str:
        return "stdio" if _clean(self.command) else "http"


@dataclass(slots=True)
class ToolsConfig:
    """Runtime tool configuration."""

    restrict_to_workspace: bool = True
    mcp_servers: dict[str, MCPServerConfig] = field(default_factory=dict)


@dataclass(slots=True)
class AuthzConfig:
    """External AuthZ service configuration."""

    enabled: bool = False
    base_url: str = ""
    request_timeout_seconds: int = 10
    outlook_mcp_url: str = ""


@dataclass(slots=True)
class BackendIdentityConfig:
    """This backend's AuthZ client identity."""

    backend_id: str = ""
    client_id: str = ""
    client_secret: str = ""
    name: str = ""
    public_base_url: str = ""


@dataclass(slots=True)
class BeaverConfig:
    """Config loaded once per backend sandbox instance."""

    agents_defaults: AgentDefaultsConfig = field(default_factory=AgentDefaultsConfig)
    providers: dict[str, ProviderConfig] = field(default_factory=dict)
    embedding: EmbeddingConfig = field(default_factory=EmbeddingConfig)
    tools: ToolsConfig = field(default_factory=ToolsConfig)
    authz: AuthzConfig = field(default_factory=AuthzConfig)
    backend_identity: BackendIdentityConfig = field(default_factory=BackendIdentityConfig)
    config_path: Path | None = None

    @property
    def default_model(self) -> str | None:
        return _clean(self.agents_defaults.model)

    @property
    def default_embedding_model(self) -> str:
        return _clean(self.embedding.model) or _clean(self.agents_defaults.embedding_model) or "text-embedding-v4"

    def resolve_provider_target(
        self,
        *,
        model: str | None = None,
        provider_name: str | None = None,
    ) -> dict[str, Any]:
        """Resolve model/provider credentials from instance config.

        Request-level model/provider overrides are allowed, but credentials are still
        read from backend config, not from Web/channel payloads.
        """

        resolved_model = _clean(model) or self.default_model
        requested_provider = _clean(provider_name)
        enabled_providers = self._enabled_provider_names()
        resolved_provider = (
            requested_provider
            if requested_provider and requested_provider in enabled_providers
            else self._infer_provider(resolved_model)
        )
        provider_cfg = self.providers.get(resolved_provider or "") if resolved_provider else None
        payload: dict[str, Any] = {
            "model": resolved_model,
            "provider_name": resolved_provider,
        }
        if provider_cfg is not None:
            payload.update(
                {
                    "api_key": provider_cfg.api_key,
                    "api_base": provider_cfg.api_base,
                    "extra_headers": dict(provider_cfg.extra_headers),
                    "request_timeout_seconds": provider_cfg.request_timeout_seconds,
                }
            )
        return {key: value for key, value in payload.items() if value not in (None, "", {})}

    def resolve_embedding_target(self) -> dict[str, Any] | None:
        """Return an explicit embedding target when configured."""

        has_explicit_embedding = any(
            [
                _clean(self.embedding.provider),
                _clean(self.embedding.api_key),
                _clean(self.embedding.api_base),
                self.embedding.extra_headers,
                self.embedding.request_timeout_seconds is not None,
            ]
        )
        if not has_explicit_embedding:
            return None

        provider_cfg = self.providers.get(_clean(self.embedding.provider) or "")
        payload: dict[str, Any] = {
            "provider": _clean(self.embedding.provider),
            "model": self.default_embedding_model,
            "api_key": _clean(self.embedding.api_key) or (provider_cfg.api_key if provider_cfg else None),
            "api_base": _clean(self.embedding.api_base) or (provider_cfg.api_base if provider_cfg else None),
            "extra_headers": self.embedding.extra_headers or (dict(provider_cfg.extra_headers) if provider_cfg else {}),
            "request_timeout_seconds": self.embedding.request_timeout_seconds
            or (provider_cfg.request_timeout_seconds if provider_cfg else None),
        }
        return {key: value for key, value in payload.items() if value not in (None, "", {})}

    def _infer_provider(self, model: str | None) -> str | None:
        configured_provider = _clean(self.agents_defaults.provider)
        if configured_provider and configured_provider != "custom":
            return configured_provider

        if model and "/" in model:
            prefix = model.split("/", 1)[0]
            if prefix in self._enabled_provider_names():
                return prefix

        enabled_providers = self._enabled_provider_names()
        if len(enabled_providers) == 1:
            return enabled_providers[0]
        return None

    def _enabled_provider_names(self) -> list[str]:
        return [
            name
            for name, provider in self.providers.items()
            if name != "custom"
            and any(
                [
                    _clean(provider.api_key),
                    _clean(provider.api_base),
                    provider.extra_headers,
                ]
            )
        ]


def _clean(value: str | None) -> str | None:
    if value is None:
        return None
    value = str(value).strip()
    return value or None