567 lines
18 KiB
Python
567 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from enum import StrEnum, auto
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shlex
|
|
import tomllib
|
|
from typing import Annotated, Any, Literal
|
|
|
|
from dotenv import dotenv_values
|
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
|
from pydantic.fields import FieldInfo
|
|
from pydantic_core import to_jsonable_python
|
|
from pydantic_settings import (
|
|
BaseSettings,
|
|
PydanticBaseSettingsSource,
|
|
SettingsConfigDict,
|
|
)
|
|
import tomli_w
|
|
|
|
from vibe.core.paths.config_paths import AGENT_DIR, CONFIG_DIR, CONFIG_FILE, PROMPT_DIR
|
|
from vibe.core.paths.global_paths import GLOBAL_ENV_FILE, SESSION_LOG_DIR
|
|
from vibe.core.prompts import SystemPrompt
|
|
from vibe.core.tools.base import BaseToolConfig
|
|
|
|
PROJECT_DOC_FILENAMES = ["AGENTS.md", "VIBE.md", ".vibe.md"]
|
|
|
|
|
|
def load_api_keys_from_env() -> None:
|
|
if GLOBAL_ENV_FILE.path.is_file():
|
|
env_vars = dotenv_values(GLOBAL_ENV_FILE.path)
|
|
for key, value in env_vars.items():
|
|
if value:
|
|
os.environ.setdefault(key, value)
|
|
|
|
|
|
class MissingAPIKeyError(RuntimeError):
|
|
def __init__(self, env_key: str, provider_name: str) -> None:
|
|
super().__init__(
|
|
f"Missing {env_key} environment variable for {provider_name} provider"
|
|
)
|
|
self.env_key = env_key
|
|
self.provider_name = provider_name
|
|
|
|
|
|
class MissingPromptFileError(RuntimeError):
|
|
def __init__(self, system_prompt_id: str, prompt_dir: str) -> None:
|
|
super().__init__(
|
|
f"Invalid system_prompt_id value: '{system_prompt_id}'. "
|
|
f"Must be one of the available prompts ({', '.join(f'{p.name.lower()}' for p in SystemPrompt)}), "
|
|
f"or correspond to a .md file in {prompt_dir}"
|
|
)
|
|
self.system_prompt_id = system_prompt_id
|
|
self.prompt_dir = prompt_dir
|
|
|
|
|
|
class WrongBackendError(RuntimeError):
|
|
def __init__(self, backend: Backend, is_mistral_api: bool) -> None:
|
|
super().__init__(
|
|
f"Wrong backend '{backend}' for {'' if is_mistral_api else 'non-'}"
|
|
f"mistral API. Use '{Backend.MISTRAL}' for mistral API and '{Backend.GENERIC}' for others."
|
|
)
|
|
self.backend = backend
|
|
self.is_mistral_api = is_mistral_api
|
|
|
|
|
|
class TomlFileSettingsSource(PydanticBaseSettingsSource):
|
|
def __init__(self, settings_cls: type[BaseSettings]) -> None:
|
|
super().__init__(settings_cls)
|
|
self.toml_data = self._load_toml()
|
|
|
|
def _load_toml(self) -> dict[str, Any]:
|
|
file = CONFIG_FILE.path
|
|
try:
|
|
with file.open("rb") as f:
|
|
return tomllib.load(f)
|
|
except FileNotFoundError:
|
|
return {}
|
|
except tomllib.TOMLDecodeError as e:
|
|
raise RuntimeError(f"Invalid TOML in {file}: {e}") from e
|
|
except OSError as e:
|
|
raise RuntimeError(f"Cannot read {file}: {e}") from e
|
|
|
|
def get_field_value(
|
|
self, field: FieldInfo, field_name: str
|
|
) -> tuple[Any, str, bool]:
|
|
return self.toml_data.get(field_name), field_name, False
|
|
|
|
def __call__(self) -> dict[str, Any]:
|
|
return self.toml_data
|
|
|
|
|
|
class ProjectContextConfig(BaseSettings):
|
|
max_chars: int = 40_000
|
|
default_commit_count: int = 5
|
|
max_doc_bytes: int = 32 * 1024
|
|
truncation_buffer: int = 1_000
|
|
max_depth: int = 3
|
|
max_files: int = 1000
|
|
max_dirs_per_level: int = 20
|
|
timeout_seconds: float = 2.0
|
|
|
|
|
|
class SessionLoggingConfig(BaseSettings):
|
|
save_dir: str = ""
|
|
session_prefix: str = "session"
|
|
enabled: bool = True
|
|
|
|
@field_validator("save_dir", mode="before")
|
|
@classmethod
|
|
def set_default_save_dir(cls, v: str) -> str:
|
|
if not v:
|
|
return str(SESSION_LOG_DIR.path)
|
|
return v
|
|
|
|
@field_validator("save_dir", mode="after")
|
|
@classmethod
|
|
def expand_save_dir(cls, v: str) -> str:
|
|
return str(Path(v).expanduser().resolve())
|
|
|
|
|
|
class Backend(StrEnum):
|
|
MISTRAL = auto()
|
|
GENERIC = auto()
|
|
|
|
|
|
class ProviderConfig(BaseModel):
|
|
name: str
|
|
api_base: str
|
|
api_key_env_var: str = ""
|
|
api_style: str = "openai"
|
|
backend: Backend = Backend.GENERIC
|
|
reasoning_field_name: str = "reasoning_content"
|
|
|
|
|
|
class _MCPBase(BaseModel):
|
|
name: str = Field(description="Short alias used to prefix tool names")
|
|
prompt: str | None = Field(
|
|
default=None, description="Optional usage hint appended to tool descriptions"
|
|
)
|
|
|
|
@field_validator("name", mode="after")
|
|
@classmethod
|
|
def normalize_name(cls, v: str) -> str:
|
|
normalized = re.sub(r"[^a-zA-Z0-9_-]", "_", v)
|
|
normalized = normalized.strip("_-")
|
|
return normalized[:256]
|
|
|
|
|
|
class _MCPHttpFields(BaseModel):
|
|
url: str = Field(description="Base URL of the MCP HTTP server")
|
|
headers: dict[str, str] = Field(
|
|
default_factory=dict,
|
|
description=(
|
|
"Additional HTTP headers when using 'http' transport (e.g., Authorization or X-API-Key)."
|
|
),
|
|
)
|
|
api_key_env: str = Field(
|
|
default="",
|
|
description=(
|
|
"Environment variable name containing an API token to send for HTTP transport."
|
|
),
|
|
)
|
|
api_key_header: str = Field(
|
|
default="Authorization",
|
|
description=(
|
|
"HTTP header name to carry the token when 'api_key_env' is set (e.g., 'Authorization' or 'X-API-Key')."
|
|
),
|
|
)
|
|
api_key_format: str = Field(
|
|
default="Bearer {token}",
|
|
description=(
|
|
"Format string for the header value when 'api_key_env' is set. Use '{token}' placeholder."
|
|
),
|
|
)
|
|
|
|
def http_headers(self) -> dict[str, str]:
|
|
hdrs = dict(self.headers or {})
|
|
env_var = (self.api_key_env or "").strip()
|
|
if env_var and (token := os.getenv(env_var)):
|
|
target = (self.api_key_header or "").strip() or "Authorization"
|
|
if not any(h.lower() == target.lower() for h in hdrs):
|
|
try:
|
|
value = (self.api_key_format or "{token}").format(token=token)
|
|
except Exception:
|
|
value = token
|
|
hdrs[target] = value
|
|
return hdrs
|
|
|
|
|
|
class MCPHttp(_MCPBase, _MCPHttpFields):
|
|
transport: Literal["http"]
|
|
|
|
|
|
class MCPStreamableHttp(_MCPBase, _MCPHttpFields):
|
|
transport: Literal["streamable-http"]
|
|
|
|
|
|
class MCPStdio(_MCPBase):
|
|
transport: Literal["stdio"]
|
|
command: str | list[str]
|
|
args: list[str] = Field(default_factory=list)
|
|
|
|
def argv(self) -> list[str]:
|
|
base = (
|
|
shlex.split(self.command)
|
|
if isinstance(self.command, str)
|
|
else list(self.command or [])
|
|
)
|
|
return [*base, *self.args] if self.args else base
|
|
|
|
|
|
MCPServer = Annotated[
|
|
MCPHttp | MCPStreamableHttp | MCPStdio, Field(discriminator="transport")
|
|
]
|
|
|
|
|
|
class ModelConfig(BaseModel):
|
|
name: str
|
|
provider: str
|
|
alias: str
|
|
temperature: float = 0.2
|
|
input_price: float = 0.0 # Price per million input tokens
|
|
output_price: float = 0.0 # Price per million output tokens
|
|
|
|
@model_validator(mode="before")
|
|
@classmethod
|
|
def _default_alias_to_name(cls, data: Any) -> Any:
|
|
if isinstance(data, dict):
|
|
if "alias" not in data or data["alias"] is None:
|
|
data["alias"] = data.get("name")
|
|
return data
|
|
|
|
|
|
DEFAULT_PROVIDERS = [
|
|
ProviderConfig(
|
|
name="mistral",
|
|
api_base="https://api.mistral.ai/v1",
|
|
api_key_env_var="MISTRAL_API_KEY",
|
|
backend=Backend.MISTRAL,
|
|
),
|
|
ProviderConfig(
|
|
name="llamacpp",
|
|
api_base="http://127.0.0.1:8080/v1",
|
|
api_key_env_var="", # NOTE: if you wish to use --api-key in llama-server, change this value
|
|
),
|
|
]
|
|
|
|
DEFAULT_MODELS = [
|
|
ModelConfig(
|
|
name="mistral-vibe-cli-latest",
|
|
provider="mistral",
|
|
alias="devstral-2",
|
|
input_price=0.4,
|
|
output_price=2.0,
|
|
),
|
|
ModelConfig(
|
|
name="devstral-small-latest",
|
|
provider="mistral",
|
|
alias="devstral-small",
|
|
input_price=0.1,
|
|
output_price=0.3,
|
|
),
|
|
ModelConfig(
|
|
name="devstral",
|
|
provider="llamacpp",
|
|
alias="local",
|
|
input_price=0.0,
|
|
output_price=0.0,
|
|
),
|
|
]
|
|
|
|
|
|
class VibeConfig(BaseSettings):
|
|
active_model: str = "devstral-2"
|
|
textual_theme: str = "terminal"
|
|
vim_keybindings: bool = False
|
|
disable_welcome_banner_animation: bool = False
|
|
displayed_workdir: str = ""
|
|
auto_compact_threshold: int = 200_000
|
|
context_warnings: bool = False
|
|
instructions: str = ""
|
|
workdir: Path | None = Field(default=None, exclude=True)
|
|
system_prompt_id: str = "cli"
|
|
include_commit_signature: bool = True
|
|
include_model_info: bool = True
|
|
include_project_context: bool = True
|
|
include_prompt_detail: bool = True
|
|
enable_update_checks: bool = True
|
|
api_timeout: float = 720.0
|
|
providers: list[ProviderConfig] = Field(
|
|
default_factory=lambda: list(DEFAULT_PROVIDERS)
|
|
)
|
|
models: list[ModelConfig] = Field(default_factory=lambda: list(DEFAULT_MODELS))
|
|
|
|
project_context: ProjectContextConfig = Field(default_factory=ProjectContextConfig)
|
|
session_logging: SessionLoggingConfig = Field(default_factory=SessionLoggingConfig)
|
|
tools: dict[str, BaseToolConfig] = Field(default_factory=dict)
|
|
tool_paths: list[Path] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"Additional directories to search for custom tools. "
|
|
"Each path may be absolute or relative to the current working directory."
|
|
),
|
|
)
|
|
|
|
mcp_servers: list[MCPServer] = Field(
|
|
default_factory=list, description="Preferred MCP server configuration entries."
|
|
)
|
|
|
|
enabled_tools: list[str] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"An explicit list of tool names/patterns to enable. If set, only these"
|
|
" tools will be active. Supports exact names, glob patterns (e.g.,"
|
|
" 'serena_*'), and regex with 're:' prefix or regex-like patterns (e.g.,"
|
|
" 're:^serena_.*' or 'serena.*')."
|
|
),
|
|
)
|
|
disabled_tools: list[str] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"A list of tool names/patterns to disable. Ignored if 'enabled_tools'"
|
|
" is set. Supports exact names, glob patterns (e.g., 'bash*'), and"
|
|
" regex with 're:' prefix or regex-like patterns."
|
|
),
|
|
)
|
|
|
|
skill_paths: list[Path] = Field(
|
|
default_factory=list,
|
|
description=(
|
|
"Additional directories to search for skills. "
|
|
"Each path may be absolute or relative to the current working directory."
|
|
),
|
|
)
|
|
|
|
model_config = SettingsConfigDict(
|
|
env_prefix="VIBE_", case_sensitive=False, extra="ignore"
|
|
)
|
|
|
|
@property
|
|
def effective_workdir(self) -> Path:
|
|
return self.workdir if self.workdir is not None else Path.cwd()
|
|
|
|
@property
|
|
def system_prompt(self) -> str:
|
|
try:
|
|
return SystemPrompt[self.system_prompt_id.upper()].read()
|
|
except KeyError:
|
|
pass
|
|
|
|
custom_sp_path = (PROMPT_DIR.path / self.system_prompt_id).with_suffix(".md")
|
|
if not custom_sp_path.is_file():
|
|
raise MissingPromptFileError(self.system_prompt_id, str(PROMPT_DIR.path))
|
|
return custom_sp_path.read_text()
|
|
|
|
def get_active_model(self) -> ModelConfig:
|
|
for model in self.models:
|
|
if model.alias == self.active_model:
|
|
return model
|
|
raise ValueError(
|
|
f"Active model '{self.active_model}' not found in configuration."
|
|
)
|
|
|
|
def get_provider_for_model(self, model: ModelConfig) -> ProviderConfig:
|
|
for provider in self.providers:
|
|
if provider.name == model.provider:
|
|
return provider
|
|
raise ValueError(
|
|
f"Provider '{model.provider}' for model '{model.name}' not found in configuration."
|
|
)
|
|
|
|
@classmethod
|
|
def settings_customise_sources(
|
|
cls,
|
|
settings_cls: type[BaseSettings],
|
|
init_settings: PydanticBaseSettingsSource,
|
|
env_settings: PydanticBaseSettingsSource,
|
|
dotenv_settings: PydanticBaseSettingsSource,
|
|
file_secret_settings: PydanticBaseSettingsSource,
|
|
) -> tuple[PydanticBaseSettingsSource, ...]:
|
|
"""Define the priority of settings sources.
|
|
|
|
Note: dotenv_settings is intentionally excluded. API keys and other
|
|
non-config environment variables are stored in .env but loaded manually
|
|
into os.environ for use by providers. Only VIBE_* prefixed environment
|
|
variables (via env_settings) and TOML config are used for Pydantic settings.
|
|
"""
|
|
return (
|
|
init_settings,
|
|
env_settings,
|
|
TomlFileSettingsSource(settings_cls),
|
|
file_secret_settings,
|
|
)
|
|
|
|
@model_validator(mode="after")
|
|
def _check_api_key(self) -> VibeConfig:
|
|
try:
|
|
active_model = self.get_active_model()
|
|
provider = self.get_provider_for_model(active_model)
|
|
api_key_env = provider.api_key_env_var
|
|
if api_key_env and not os.getenv(api_key_env):
|
|
raise MissingAPIKeyError(api_key_env, provider.name)
|
|
except ValueError:
|
|
pass
|
|
return self
|
|
|
|
@model_validator(mode="after")
|
|
def _check_api_backend_compatibility(self) -> VibeConfig:
|
|
try:
|
|
active_model = self.get_active_model()
|
|
provider = self.get_provider_for_model(active_model)
|
|
MISTRAL_API_BASES = [
|
|
"https://codestral.mistral.ai",
|
|
"https://api.mistral.ai",
|
|
]
|
|
is_mistral_api = any(
|
|
provider.api_base.startswith(api_base) for api_base in MISTRAL_API_BASES
|
|
)
|
|
if (is_mistral_api and provider.backend != Backend.MISTRAL) or (
|
|
not is_mistral_api and provider.backend != Backend.GENERIC
|
|
):
|
|
raise WrongBackendError(provider.backend, is_mistral_api)
|
|
|
|
except ValueError:
|
|
pass
|
|
return self
|
|
|
|
@field_validator("tool_paths", mode="before")
|
|
@classmethod
|
|
def _expand_tool_paths(cls, v: Any) -> list[Path]:
|
|
if not v:
|
|
return []
|
|
return [Path(p).expanduser().resolve() for p in v]
|
|
|
|
@field_validator("skill_paths", mode="before")
|
|
@classmethod
|
|
def _expand_skill_paths(cls, v: Any) -> list[Path]:
|
|
if not v:
|
|
return []
|
|
return [Path(p).expanduser().resolve() for p in v]
|
|
|
|
@field_validator("workdir", mode="before")
|
|
@classmethod
|
|
def _expand_workdir(cls, v: Any) -> Path | None:
|
|
if v is None or (isinstance(v, str) and not v.strip()):
|
|
return None
|
|
|
|
if isinstance(v, str):
|
|
v = Path(v).expanduser().resolve()
|
|
elif isinstance(v, Path):
|
|
v = v.expanduser().resolve()
|
|
if not v.is_dir():
|
|
raise ValueError(
|
|
f"Tried to set {v} as working directory, path doesn't exist"
|
|
)
|
|
return v
|
|
|
|
@field_validator("tools", mode="before")
|
|
@classmethod
|
|
def _normalize_tool_configs(cls, v: Any) -> dict[str, BaseToolConfig]:
|
|
if not isinstance(v, dict):
|
|
return {}
|
|
|
|
normalized: dict[str, BaseToolConfig] = {}
|
|
for tool_name, tool_config in v.items():
|
|
if isinstance(tool_config, BaseToolConfig):
|
|
normalized[tool_name] = tool_config
|
|
elif isinstance(tool_config, dict):
|
|
normalized[tool_name] = BaseToolConfig.model_validate(tool_config)
|
|
else:
|
|
normalized[tool_name] = BaseToolConfig()
|
|
|
|
return normalized
|
|
|
|
@model_validator(mode="after")
|
|
def _validate_model_uniqueness(self) -> VibeConfig:
|
|
seen_aliases: set[str] = set()
|
|
for model in self.models:
|
|
if model.alias in seen_aliases:
|
|
raise ValueError(
|
|
f"Duplicate model alias found: '{model.alias}'. Aliases must be unique."
|
|
)
|
|
seen_aliases.add(model.alias)
|
|
return self
|
|
|
|
@model_validator(mode="after")
|
|
def _check_system_prompt(self) -> VibeConfig:
|
|
_ = self.system_prompt
|
|
return self
|
|
|
|
@classmethod
|
|
def save_updates(cls, updates: dict[str, Any]) -> None:
|
|
CONFIG_DIR.path.mkdir(parents=True, exist_ok=True)
|
|
current_config = TomlFileSettingsSource(cls).toml_data
|
|
|
|
def deep_merge(target: dict, source: dict) -> None:
|
|
for key, value in source.items():
|
|
if (
|
|
key in target
|
|
and isinstance(target.get(key), dict)
|
|
and isinstance(value, dict)
|
|
):
|
|
deep_merge(target[key], value)
|
|
elif (
|
|
key in target
|
|
and isinstance(target.get(key), list)
|
|
and isinstance(value, list)
|
|
):
|
|
if key in {"providers", "models"}:
|
|
target[key] = value
|
|
else:
|
|
target[key] = list(set(value + target[key]))
|
|
else:
|
|
target[key] = value
|
|
|
|
deep_merge(current_config, updates)
|
|
cls.dump_config(
|
|
to_jsonable_python(current_config, exclude_none=True, fallback=str)
|
|
)
|
|
|
|
@classmethod
|
|
def dump_config(cls, config: dict[str, Any]) -> None:
|
|
with CONFIG_FILE.path.open("wb") as f:
|
|
tomli_w.dump(config, f)
|
|
|
|
@classmethod
|
|
def _get_agent_config(cls, agent: str | None) -> dict[str, Any] | None:
|
|
if agent is None:
|
|
return None
|
|
|
|
agent_config_path = (AGENT_DIR.path / agent).with_suffix(".toml")
|
|
try:
|
|
return tomllib.load(agent_config_path.open("rb"))
|
|
except FileNotFoundError:
|
|
raise ValueError(
|
|
f"Config '{agent}.toml' for agent not found in {AGENT_DIR.path}"
|
|
)
|
|
|
|
@classmethod
|
|
def _migrate(cls) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def load(cls, agent: str | None = None, **overrides: Any) -> VibeConfig:
|
|
cls._migrate()
|
|
agent_config = cls._get_agent_config(agent)
|
|
init_data = {**(agent_config or {}), **overrides}
|
|
return cls(**init_data)
|
|
|
|
@classmethod
|
|
def create_default(cls) -> dict[str, Any]:
|
|
try:
|
|
config = cls()
|
|
except MissingAPIKeyError:
|
|
config = cls.model_construct()
|
|
|
|
config_dict = config.model_dump(mode="json", exclude_none=True)
|
|
|
|
from vibe.core.tools.manager import ToolManager
|
|
|
|
tool_defaults = ToolManager.discover_tool_defaults()
|
|
if tool_defaults:
|
|
config_dict["tools"] = tool_defaults
|
|
|
|
return config_dict
|