codemaster-cli 2.2.0__py3-none-any.whl

vibe/core/llm/format.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+from vibe.core.tools.base import BaseTool
+from vibe.core.types import (
+    AvailableFunction,
+    AvailableTool,
+    LLMMessage,
+    Role,
+    StrToolChoice,
+)
+
+if TYPE_CHECKING:
+    from vibe.core.tools.manager import ToolManager
+
+
+class ParsedToolCall(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    tool_name: str
+    raw_args: dict[str, Any]
+    call_id: str = ""
+
+
+class ResolvedToolCall(BaseModel):
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+    tool_name: str
+    tool_class: type[BaseTool]
+    validated_args: BaseModel
+    call_id: str = ""
+
+    @property
+    def args_dict(self) -> dict[str, Any]:
+        return self.validated_args.model_dump()
+
+
+class FailedToolCall(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    tool_name: str
+    call_id: str
+    error: str
+
+
+class ParsedMessage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    tool_calls: list[ParsedToolCall]
+
+
+class ResolvedMessage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    tool_calls: list[ResolvedToolCall]
+    failed_calls: list[FailedToolCall] = Field(default_factory=list)
+
+
+class APIToolFormatHandler:
+    @property
+    def name(self) -> str:
+        return "api"
+
+    def get_available_tools(self, tool_manager: ToolManager) -> list[AvailableTool]:
+        return [
+            AvailableTool(
+                function=AvailableFunction(
+                    name=tool_class.get_name(),
+                    description=tool_class.description,
+                    parameters=tool_class.get_parameters(),
+                )
+            )
+            for tool_class in tool_manager.available_tools.values()
+        ]
+
+    def get_tool_choice(self) -> StrToolChoice | AvailableTool:
+        return "auto"
+
+    def process_api_response_message(self, message: Any) -> LLMMessage:
+        clean_message = {
+            "role": message.role,
+            "content": message.content,
+            "reasoning_content": getattr(message, "reasoning_content", None),
+            "reasoning_signature": getattr(message, "reasoning_signature", None),
+        }
+
+        if message.tool_calls:
+            clean_message["tool_calls"] = [
+                {
+                    "id": tc.id,
+                    "index": tc.index,
+                    "type": "function",
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments,
+                    },
+                }
+                for tc in message.tool_calls
+            ]
+
+        return LLMMessage.model_validate(clean_message)
+
+    def parse_message(self, message: LLMMessage) -> ParsedMessage:
+        tool_calls = []
+
+        api_tool_calls = message.tool_calls or []
+        for tc in api_tool_calls:
+            if not (function_call := tc.function):
+                continue
+            try:
+                args = json.loads(function_call.arguments or "{}")
+            except json.JSONDecodeError:
+                args = {}
+
+            tool_calls.append(
+                ParsedToolCall(
+                    tool_name=function_call.name or "",
+                    raw_args=args,
+                    call_id=tc.id or "",
+                )
+            )
+
+        return ParsedMessage(tool_calls=tool_calls)
+
+    def resolve_tool_calls(
+        self, parsed: ParsedMessage, tool_manager: ToolManager
+    ) -> ResolvedMessage:
+        resolved_calls = []
+        failed_calls = []
+
+        active_tools = tool_manager.available_tools
+
+        for parsed_call in parsed.tool_calls:
+            tool_class = active_tools.get(parsed_call.tool_name)
+            if not tool_class:
+                failed_calls.append(
+                    FailedToolCall(
+                        tool_name=parsed_call.tool_name,
+                        call_id=parsed_call.call_id,
+                        error=f"Unknown tool '{parsed_call.tool_name}'",
+                    )
+                )
+                continue
+
+            args_model, _ = tool_class._get_tool_args_results()
+            try:
+                validated_args = args_model.model_validate(parsed_call.raw_args)
+                resolved_calls.append(
+                    ResolvedToolCall(
+                        tool_name=parsed_call.tool_name,
+                        tool_class=tool_class,
+                        validated_args=validated_args,
+                        call_id=parsed_call.call_id,
+                    )
+                )
+            except ValidationError as e:
+                failed_calls.append(
+                    FailedToolCall(
+                        tool_name=parsed_call.tool_name,
+                        call_id=parsed_call.call_id,
+                        error=f"Invalid arguments: {e}",
+                    )
+                )
+
+        return ResolvedMessage(tool_calls=resolved_calls, failed_calls=failed_calls)
+
+    def create_tool_response_message(
+        self, tool_call: ResolvedToolCall, result_text: str
+    ) -> LLMMessage:
+        return LLMMessage(
+            role=Role.tool,
+            tool_call_id=tool_call.call_id,
+            name=tool_call.tool_name,
+            content=result_text,
+        )
+
+    def create_failed_tool_response_message(
+        self, failed: FailedToolCall, error_content: str
+    ) -> LLMMessage:
+        return LLMMessage(
+            role=Role.tool,
+            tool_call_id=failed.call_id,
+            name=failed.tool_name,
+            content=error_content,
+        )

vibe/core/llm/message_utils.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from vibe.core.types import LLMMessage, Role
+
+
+def merge_consecutive_user_messages(messages: list[LLMMessage]) -> list[LLMMessage]:
+    """Merge consecutive user messages into a single message.
+
+    This handles cases where middleware injects messages resulting in
+    consecutive user messages before sending to the API.
+    """
+    result: list[LLMMessage] = []
+    for msg in messages:
+        if result and result[-1].role == Role.user and msg.role == Role.user:
+            prev_content = result[-1].content or ""
+            curr_content = msg.content or ""
+            merged_content = f"{prev_content}\n\n{curr_content}".strip()
+            result[-1] = LLMMessage(
+                role=Role.user, content=merged_content, message_id=result[-1].message_id
+            )
+        else:
+            result.append(msg)
+
+    return result

vibe/core/llm/types.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+import types
+from typing import TYPE_CHECKING, Protocol
+
+from vibe.core.types import AvailableTool, LLMChunk, LLMMessage, StrToolChoice
+
+if TYPE_CHECKING:
+    from vibe.core.config import ModelConfig
+
+
+class BackendLike(Protocol):
+    """Port protocol for dependency-injectable LLM backends.
+
+    Any backend used by AgentLoop should implement this async context manager
+    interface with `complete`, `complete_streaming` and `count_tokens` methods.
+    """
+
+    async def __aenter__(self) -> BackendLike: ...
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None: ...
+
+    async def complete(
+        self,
+        *,
+        model: ModelConfig,
+        messages: list[LLMMessage],
+        temperature: float,
+        tools: list[AvailableTool] | None,
+        max_tokens: int | None,
+        tool_choice: StrToolChoice | AvailableTool | None,
+        extra_headers: dict[str, str] | None,
+    ) -> LLMChunk:
+        """Complete a chat conversation using the specified model and provider.
+
+        Args:
+            model: Model configuration
+            messages: List of conversation messages
+            temperature: Sampling temperature (0.0 to 1.0)
+            tools: Optional list of available tools
+            max_tokens: Maximum tokens to generate
+            tool_choice: How to choose tools (auto, none, or specific tool)
+            extra_headers: Additional HTTP headers to include
+
+        Returns:
+            LLMChunk containing the response message and usage information
+
+        Raises:
+            BackendError: If the API request fails
+        """
+        ...
+
+    # Note: actual implementation should be an async function,
+    # but we can't make this one async, as it would lead to wrong type inference
+    # https://stackoverflow.com/a/68911014
+    def complete_streaming(
+        self,
+        *,
+        model: ModelConfig,
+        messages: list[LLMMessage],
+        temperature: float,
+        tools: list[AvailableTool] | None,
+        max_tokens: int | None,
+        tool_choice: StrToolChoice | AvailableTool | None,
+        extra_headers: dict[str, str] | None,
+    ) -> AsyncGenerator[LLMChunk, None]:
+        """Equivalent of the complete method, but yields LLMEvent objects
+        instead of a single LLMEvent.
+
+        Args:
+            model: Model configuration
+            messages: List of conversation messages
+            temperature: Sampling temperature (0.0 to 1.0)
+            tools: Optional list of available tools
+            max_tokens: Maximum tokens to generate
+            tool_choice: How to choose tools (auto, none, or specific tool)
+            extra_headers: Additional HTTP headers to include
+
+        Returns:
+            AsyncGenerator[LLMEvent, None] yielding LLMEvent objects
+
+        Raises:
+            BackendError: If the API request fails
+        """
+        ...
+
+    async def count_tokens(
+        self,
+        *,
+        model: ModelConfig,
+        messages: list[LLMMessage],
+        temperature: float = 0.0,
+        tools: list[AvailableTool] | None,
+        tool_choice: StrToolChoice | AvailableTool | None = None,
+        extra_headers: dict[str, str] | None,
+    ) -> int:
+        """Count the number of tokens in the prompt without generating a real response.
+
+        This is useful for:
+        - Determining system prompt token count
+        - Checking context size after compaction
+        - Pre-flight token validation
+
+        Args:
+            model: Model configuration
+            messages: List of messages to count tokens for
+            temperature: Sampling temperature
+            tools: Optional list of available tools
+            tool_choice: How to choose tools
+            extra_headers: Additional HTTP headers to include
+
+        Returns:
+            The number of prompt tokens
+        """
+        ...

vibe/core/middleware.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import StrEnum, auto
+from typing import TYPE_CHECKING, Any, Protocol
+
+from vibe.core.agents import AgentProfile
+from vibe.core.agents.models import BuiltinAgentName
+from vibe.core.utils import VIBE_WARNING_TAG
+
+if TYPE_CHECKING:
+    from vibe.core.config import VibeConfig
+    from vibe.core.types import AgentStats, LLMMessage
+
+
+class MiddlewareAction(StrEnum):
+    CONTINUE = auto()
+    STOP = auto()
+    COMPACT = auto()
+    INJECT_MESSAGE = auto()
+
+
+class ResetReason(StrEnum):
+    STOP = auto()
+    COMPACT = auto()
+
+
+@dataclass
+class ConversationContext:
+    messages: list[LLMMessage]
+    stats: AgentStats
+    config: VibeConfig
+
+
+@dataclass
+class MiddlewareResult:
+    action: MiddlewareAction = MiddlewareAction.CONTINUE
+    message: str | None = None
+    reason: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class ConversationMiddleware(Protocol):
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult: ...
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None: ...
+
+
+class TurnLimitMiddleware:
+    def __init__(self, max_turns: int) -> None:
+        self.max_turns = max_turns
+
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        if context.stats.steps - 1 >= self.max_turns:
+            return MiddlewareResult(
+                action=MiddlewareAction.STOP,
+                reason=f"Turn limit of {self.max_turns} reached",
+            )
+        return MiddlewareResult()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        pass
+
+
+class PriceLimitMiddleware:
+    def __init__(self, max_price: float) -> None:
+        self.max_price = max_price
+
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        if context.stats.session_cost > self.max_price:
+            return MiddlewareResult(
+                action=MiddlewareAction.STOP,
+                reason=f"Price limit exceeded: ${context.stats.session_cost:.4f} > ${self.max_price:.2f}",
+            )
+        return MiddlewareResult()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        pass
+
+
+class AutoCompactMiddleware:
+    def __init__(self, threshold: int) -> None:
+        self.threshold = threshold
+
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        if context.stats.context_tokens >= self.threshold:
+            return MiddlewareResult(
+                action=MiddlewareAction.COMPACT,
+                metadata={
+                    "old_tokens": context.stats.context_tokens,
+                    "threshold": self.threshold,
+                },
+            )
+        return MiddlewareResult()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        pass
+
+
+class ContextWarningMiddleware:
+    def __init__(
+        self, threshold_percent: float = 0.5, max_context: int | None = None
+    ) -> None:
+        self.threshold_percent = threshold_percent
+        self.max_context = max_context
+        self.has_warned = False
+
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        if self.has_warned:
+            return MiddlewareResult()
+
+        max_context = self.max_context
+        if max_context is None:
+            return MiddlewareResult()
+
+        if context.stats.context_tokens >= max_context * self.threshold_percent:
+            self.has_warned = True
+
+            percentage_used = (context.stats.context_tokens / max_context) * 100
+            warning_msg = f"<{VIBE_WARNING_TAG}>You have used {percentage_used:.0f}% of your total context ({context.stats.context_tokens:,}/{max_context:,} tokens)</{VIBE_WARNING_TAG}>"
+
+            return MiddlewareResult(
+                action=MiddlewareAction.INJECT_MESSAGE, message=warning_msg
+            )
+
+        return MiddlewareResult()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        self.has_warned = False
+
+
+PLAN_AGENT_REMINDER = f"""<{VIBE_WARNING_TAG}>Plan mode is active. The user indicated that they do not want you to execute yet -- you MUST NOT make any edits, run any non-readonly tools (including changing configs or making commits), or otherwise make any changes to the system. This supersedes any other instructions you have received (for example, to make edits). Instead, you should:
+1. Answer the user's query comprehensively
+2. When you're done researching, present your plan by giving the full plan and not doing further tool calls to return input to the user. Do NOT make any file changes or run any tools that modify the system state in any way until the user has confirmed the plan.</{VIBE_WARNING_TAG}>"""
+
+PLAN_AGENT_EXIT = f"""<{VIBE_WARNING_TAG}>Plan mode has ended. If you have a plan ready, you can now start executing it. If not, you can now use editing tools and make changes to the system.</{VIBE_WARNING_TAG}>"""
+
+
+class PlanAgentMiddleware:
+    def __init__(
+        self,
+        profile_getter: Callable[[], AgentProfile],
+        reminder: str = PLAN_AGENT_REMINDER,
+        exit_message: str = PLAN_AGENT_EXIT,
+    ) -> None:
+        self._profile_getter = profile_getter
+        self.reminder = reminder
+        self.exit_message = exit_message
+        self._was_plan_agent = False
+
+    def _is_plan_agent(self) -> bool:
+        return self._profile_getter().name == BuiltinAgentName.PLAN
+
+    async def before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        is_plan = self._is_plan_agent()
+        was_plan = self._was_plan_agent
+
+        if was_plan and not is_plan:
+            self._was_plan_agent = False
+            return MiddlewareResult(
+                action=MiddlewareAction.INJECT_MESSAGE, message=self.exit_message
+            )
+
+        if is_plan and not was_plan:
+            self._was_plan_agent = True
+            return MiddlewareResult(
+                action=MiddlewareAction.INJECT_MESSAGE, message=self.reminder
+            )
+
+        self._was_plan_agent = is_plan
+
+        return MiddlewareResult()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        self._was_plan_agent = False
+
+
+class MiddlewarePipeline:
+    def __init__(self) -> None:
+        self.middlewares: list[ConversationMiddleware] = []
+
+    def add(self, middleware: ConversationMiddleware) -> MiddlewarePipeline:
+        self.middlewares.append(middleware)
+        return self
+
+    def clear(self) -> None:
+        self.middlewares.clear()
+
+    def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
+        for mw in self.middlewares:
+            mw.reset(reset_reason)
+
+    async def run_before_turn(self, context: ConversationContext) -> MiddlewareResult:
+        messages_to_inject = []
+
+        for mw in self.middlewares:
+            result = await mw.before_turn(context)
+            if result.action == MiddlewareAction.INJECT_MESSAGE and result.message:
+                messages_to_inject.append(result.message)
+            elif result.action in {MiddlewareAction.STOP, MiddlewareAction.COMPACT}:
+                return result
+        if messages_to_inject:
+            combined_message = "\n\n".join(messages_to_inject)
+            return MiddlewareResult(
+                action=MiddlewareAction.INJECT_MESSAGE, message=combined_message
+            )
+
+        return MiddlewareResult()

vibe/core/output_formatters.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import json
+import sys
+from typing import TextIO
+
+from vibe.core.types import AssistantEvent, BaseEvent, LLMMessage, OutputFormat
+
+
+class OutputFormatter(ABC):
+    def __init__(self, stream: TextIO = sys.stdout) -> None:
+        self.stream = stream
+        self._messages: list[LLMMessage] = []
+        self._final_response: str | None = None
+
+    @abstractmethod
+    def on_message_added(self, message: LLMMessage) -> None:
+        pass
+
+    @abstractmethod
+    def on_event(self, event: BaseEvent) -> None:
+        pass
+
+    @abstractmethod
+    def finalize(self) -> str | None:
+        """Finalize output and return any final text to be printed.
+
+        Returns:
+            String to print, or None if formatter handles its own output
+        """
+        pass
+
+
+class TextOutputFormatter(OutputFormatter):
+    def on_message_added(self, message: LLMMessage) -> None:
+        self._messages.append(message)
+
+    def on_event(self, event: BaseEvent) -> None:
+        if isinstance(event, AssistantEvent):
+            self._final_response = event.content
+
+    def finalize(self) -> str | None:
+        return self._final_response
+
+
+class JsonOutputFormatter(OutputFormatter):
+    def on_message_added(self, message: LLMMessage) -> None:
+        self._messages.append(message)
+
+    def on_event(self, event: BaseEvent) -> None:
+        pass
+
+    def finalize(self) -> str | None:
+        messages_data = [msg.model_dump(mode="json") for msg in self._messages]
+        json.dump(messages_data, self.stream, indent=2, ensure_ascii=False)
+        self.stream.write("\n")
+        self.stream.flush()
+        return None
+
+
+class StreamingJsonOutputFormatter(OutputFormatter):
+    def on_message_added(self, message: LLMMessage) -> None:
+        json.dump(message.model_dump(mode="json"), self.stream, ensure_ascii=False)
+        self.stream.write("\n")
+        self.stream.flush()
+
+    def on_event(self, event: BaseEvent) -> None:
+        pass
+
+    def finalize(self) -> str | None:
+        return None
+
+
+def create_formatter(
+    format_type: OutputFormat, stream: TextIO = sys.stdout
+) -> OutputFormatter:
+    formatters = {
+        OutputFormat.TEXT: TextOutputFormatter,
+        OutputFormat.JSON: JsonOutputFormatter,
+        OutputFormat.STREAMING: StreamingJsonOutputFormatter,
+    }
+
+    formatter_class = formatters.get(format_type, TextOutputFormatter)
+    return formatter_class(stream)

vibe/core/paths/config_paths.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+
+from vibe.core.paths.global_paths import VIBE_HOME, GlobalPath
+from vibe.core.trusted_folders import trusted_folders_manager
+
+_config_paths_locked: bool = True
+
+
+class ConfigPath(GlobalPath):
+    @property
+    def path(self) -> Path:
+        if _config_paths_locked:
+            raise RuntimeError("Config path is locked")
+        return super().path
+
+
+def _resolve_config_path(basename: str, type: Literal["file", "dir"]) -> Path:
+    cwd = Path.cwd()
+    is_folder_trusted = trusted_folders_manager.is_trusted(cwd)
+    if not is_folder_trusted:
+        return VIBE_HOME.path / basename
+    if type == "file":
+        if (candidate := cwd / ".vibe" / basename).is_file():
+            return candidate
+    elif type == "dir":
+        if (candidate := cwd / ".vibe" / basename).is_dir():
+            return candidate
+    return VIBE_HOME.path / basename
+
+
+def resolve_local_tools_dir(dir: Path) -> Path | None:
+    if not trusted_folders_manager.is_trusted(dir):
+        return None
+    if (candidate := dir / ".vibe" / "tools").is_dir():
+        return candidate
+    return None
+
+
+def resolve_local_skills_dirs(dir: Path) -> list[Path]:
+    if not trusted_folders_manager.is_trusted(dir):
+        return []
+    return [
+        candidate
+        for candidate in [dir / ".vibe" / "skills", dir / ".agents" / "skills"]
+        if candidate.is_dir()
+    ]
+
+
+def resolve_local_agents_dir(dir: Path) -> Path | None:
+    if not trusted_folders_manager.is_trusted(dir):
+        return None
+    if (candidate := dir / ".vibe" / "agents").is_dir():
+        return candidate
+    return None
+
+
+def unlock_config_paths() -> None:
+    global _config_paths_locked
+    _config_paths_locked = False
+
+
+CONFIG_FILE = ConfigPath(lambda: _resolve_config_path("config.toml", "file"))
+CONFIG_DIR = ConfigPath(lambda: CONFIG_FILE.path.parent)
+PROMPTS_DIR = ConfigPath(lambda: _resolve_config_path("prompts", "dir"))
+HISTORY_FILE = ConfigPath(lambda: _resolve_config_path("vibehistory", "file"))