Initial commit
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Laure Hugo <laure.hugo@mistral.ai> Co-Authored-By: Benjamin Trom <benjamin.trom@mistral.ai> Co-Authored-By: Mathias Gesbert <mathias.gesbert@ext.mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Valentin Berard <val@mistral.ai> Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
commit
fa15fc977b
200 changed files with 30484 additions and 0 deletions
5
vibe/__init__.py
Normal file
5
vibe/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
VIBE_ROOT = Path(__file__).parent
|
||||
0
vibe/acp/__init__.py
Normal file
0
vibe/acp/__init__.py
Normal file
441
vibe/acp/acp_agent.py
Normal file
441
vibe/acp/acp_agent.py
Normal file
|
|
@ -0,0 +1,441 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Any, cast, override
|
||||
|
||||
from acp import (
|
||||
PROTOCOL_VERSION,
|
||||
Agent as AcpAgent,
|
||||
AgentSideConnection,
|
||||
AuthenticateRequest,
|
||||
CancelNotification,
|
||||
InitializeRequest,
|
||||
InitializeResponse,
|
||||
LoadSessionRequest,
|
||||
NewSessionRequest,
|
||||
NewSessionResponse,
|
||||
PromptRequest,
|
||||
PromptResponse,
|
||||
RequestError,
|
||||
RequestPermissionRequest,
|
||||
SessionNotification,
|
||||
SetSessionModelRequest,
|
||||
SetSessionModelResponse,
|
||||
SetSessionModeRequest,
|
||||
SetSessionModeResponse,
|
||||
stdio_streams,
|
||||
)
|
||||
from acp.helpers import ContentBlock, SessionUpdate
|
||||
from acp.schema import (
|
||||
AgentCapabilities,
|
||||
AgentMessageChunk,
|
||||
AllowedOutcome,
|
||||
AuthenticateResponse,
|
||||
AuthMethod,
|
||||
Implementation,
|
||||
ModelInfo,
|
||||
PromptCapabilities,
|
||||
SessionModelState,
|
||||
SessionModeState,
|
||||
TextContentBlock,
|
||||
TextResourceContents,
|
||||
ToolCall,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import BaseAcpTool
|
||||
from vibe.acp.tools.session_update import (
|
||||
tool_call_session_update,
|
||||
tool_result_session_update,
|
||||
)
|
||||
from vibe.acp.utils import TOOL_OPTIONS, ToolOption, VibeSessionMode
|
||||
from vibe.core import __version__
|
||||
from vibe.core.agent import Agent as VibeAgent
|
||||
from vibe.core.autocompletion.path_prompt_adapter import render_path_prompt
|
||||
from vibe.core.config import MissingAPIKeyError, VibeConfig, load_api_keys_from_env
|
||||
from vibe.core.types import (
|
||||
AssistantEvent,
|
||||
AsyncApprovalCallback,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
)
|
||||
from vibe.core.utils import CancellationReason, get_user_cancellation_message
|
||||
|
||||
|
||||
class AcpSession(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
id: str
|
||||
agent: VibeAgent
|
||||
mode_id: VibeSessionMode = VibeSessionMode.APPROVAL_REQUIRED
|
||||
task: asyncio.Task[None] | None = None
|
||||
|
||||
|
||||
class VibeAcpAgent(AcpAgent):
|
||||
def __init__(self, connection: AgentSideConnection) -> None:
|
||||
self.sessions: dict[str, AcpSession] = {}
|
||||
self.connection = connection
|
||||
self.client_capabilities = None
|
||||
|
||||
@override
|
||||
async def initialize(self, params: InitializeRequest) -> InitializeResponse:
|
||||
self.client_capabilities = params.clientCapabilities
|
||||
|
||||
# The ACP Agent process can be launched in 3 different ways, depending on installation
|
||||
# - dev mode: `uv run vibe-acp`, ran from the project root
|
||||
# - uv tool install: `vibe-acp`, similar to dev mode, but uv takes care of path resolution
|
||||
# - bundled binary: `./vibe-acp` from binary location
|
||||
# The 2 first modes are working similarly, under the hood uv runs `/some/python /my/entrypoint.py``
|
||||
# The last mode is quite different as our bundler also includes the python install.
|
||||
# So sys.executable is already /path/to/binary/vibe-acp.
|
||||
# For this reason, we make a distinction in the way we call the setup command
|
||||
command = sys.executable
|
||||
if "python" not in Path(command).name:
|
||||
# It's the case for bundled binaries, we don't need any other arguments
|
||||
args = ["--setup"]
|
||||
else:
|
||||
script_name = sys.argv[0]
|
||||
args = [script_name, "--setup"]
|
||||
|
||||
auth_methods = [
|
||||
AuthMethod(
|
||||
id="vibe-setup",
|
||||
name="Register your API Key",
|
||||
description="Register your API Key inside Mistral Vibe",
|
||||
field_meta={
|
||||
"terminal-auth": {
|
||||
"command": command,
|
||||
"args": args,
|
||||
"label": "Mistral Vibe Setup",
|
||||
}
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
response = InitializeResponse(
|
||||
agentCapabilities=AgentCapabilities(
|
||||
loadSession=False,
|
||||
promptCapabilities=PromptCapabilities(
|
||||
audio=False, embeddedContext=True, image=False
|
||||
),
|
||||
),
|
||||
protocolVersion=PROTOCOL_VERSION,
|
||||
agentInfo=Implementation(
|
||||
name="@mistralai/mistral-vibe",
|
||||
title="Mistral Vibe",
|
||||
version=__version__,
|
||||
),
|
||||
authMethods=auth_methods,
|
||||
)
|
||||
return response
|
||||
|
||||
@override
|
||||
async def authenticate(
|
||||
self, params: AuthenticateRequest
|
||||
) -> AuthenticateResponse | None:
|
||||
raise NotImplementedError("Not implemented yet")
|
||||
|
||||
@override
|
||||
async def newSession(self, params: NewSessionRequest) -> NewSessionResponse:
|
||||
capability_disabled_tools = self._get_disabled_tools_from_capabilities()
|
||||
load_api_keys_from_env()
|
||||
try:
|
||||
config = VibeConfig.load(
|
||||
workdir=Path(params.cwd),
|
||||
tool_paths=[str(VIBE_ROOT / "acp" / "tools" / "builtins")],
|
||||
disabled_tools=capability_disabled_tools,
|
||||
)
|
||||
except MissingAPIKeyError as e:
|
||||
raise RequestError.auth_required({
|
||||
"message": "You must be authenticated before creating a new session"
|
||||
}) from e
|
||||
|
||||
agent = VibeAgent(config=config, auto_approve=False, enable_streaming=True)
|
||||
# NOTE: For now, we pin session.id to agent.session_id right after init time.
|
||||
# We should just use agent.session_id everywhere, but it can still change during
|
||||
# session lifetime (e.g. agent.compact is called).
|
||||
# We should refactor agent.session_id to make it immutable in ACP context.
|
||||
session = AcpSession(id=agent.session_id, agent=agent)
|
||||
self.sessions[session.id] = session
|
||||
|
||||
if not agent.auto_approve:
|
||||
agent.set_approval_callback(
|
||||
self._create_approval_callback(agent.session_id)
|
||||
)
|
||||
|
||||
response = NewSessionResponse(
|
||||
sessionId=agent.session_id,
|
||||
models=SessionModelState(
|
||||
currentModelId=agent.config.active_model,
|
||||
availableModels=[
|
||||
ModelInfo(modelId=model.alias, name=model.alias)
|
||||
for model in agent.config.models
|
||||
],
|
||||
),
|
||||
modes=SessionModeState(
|
||||
currentModeId=session.mode_id,
|
||||
availableModes=VibeSessionMode.get_all_acp_session_modes(),
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
||||
def _get_disabled_tools_from_capabilities(self) -> list[str]:
|
||||
if not self.client_capabilities:
|
||||
return []
|
||||
|
||||
disabled: list[str] = []
|
||||
|
||||
if not self.client_capabilities.terminal:
|
||||
disabled.append("bash")
|
||||
|
||||
if fs := self.client_capabilities.fs:
|
||||
if not fs.readTextFile:
|
||||
disabled.append("read_file")
|
||||
if not fs.writeTextFile:
|
||||
disabled.append("write_file")
|
||||
disabled.append("search_replace")
|
||||
|
||||
return disabled
|
||||
|
||||
def _create_approval_callback(self, session_id: str) -> AsyncApprovalCallback:
|
||||
|
||||
async def approval_callback(
|
||||
tool_name: str, args: dict[str, Any], tool_call_id: str
|
||||
) -> tuple[str, str | None]:
|
||||
# Create the tool call update
|
||||
tool_call = ToolCall(toolCallId=tool_call_id)
|
||||
|
||||
# Request permission from the user
|
||||
request = RequestPermissionRequest(
|
||||
sessionId=session_id, toolCall=tool_call, options=TOOL_OPTIONS
|
||||
)
|
||||
|
||||
response = await self.connection.requestPermission(request)
|
||||
|
||||
# Parse the response using isinstance for proper type narrowing
|
||||
if response.outcome.outcome == "selected":
|
||||
outcome = cast(AllowedOutcome, response.outcome)
|
||||
return self._handle_permission_selection(outcome.optionId)
|
||||
else:
|
||||
return (
|
||||
"n",
|
||||
str(
|
||||
get_user_cancellation_message(
|
||||
CancellationReason.OPERATION_CANCELLED
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
return approval_callback
|
||||
|
||||
@staticmethod
|
||||
def _handle_permission_selection(option_id: str) -> tuple[str, str | None]:
|
||||
match option_id:
|
||||
case ToolOption.ALLOW_ONCE:
|
||||
return ("y", None)
|
||||
case ToolOption.ALLOW_ALWAYS:
|
||||
return ("a", None)
|
||||
case ToolOption.REJECT_ONCE:
|
||||
return ("n", "User rejected the tool call, provide an alternative plan")
|
||||
case _:
|
||||
return ("n", f"Unknown option: {option_id}")
|
||||
|
||||
def _get_session(self, session_id: str) -> AcpSession:
|
||||
if session_id not in self.sessions:
|
||||
raise RequestError.invalid_params({"session": "Not found"})
|
||||
return self.sessions[session_id]
|
||||
|
||||
@override
|
||||
async def loadSession(self, params: LoadSessionRequest) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
@override
|
||||
async def setSessionMode(
|
||||
self, params: SetSessionModeRequest
|
||||
) -> SetSessionModeResponse | None:
|
||||
session = self._get_session(params.sessionId)
|
||||
|
||||
if not VibeSessionMode.is_valid(params.modeId):
|
||||
return None
|
||||
|
||||
session.mode_id = VibeSessionMode(params.modeId)
|
||||
session.agent.auto_approve = params.modeId == VibeSessionMode.AUTO_APPROVE
|
||||
|
||||
return SetSessionModeResponse()
|
||||
|
||||
@override
|
||||
async def setSessionModel(
|
||||
self, params: SetSessionModelRequest
|
||||
) -> SetSessionModelResponse | None:
|
||||
session = self._get_session(params.sessionId)
|
||||
|
||||
model_aliases = [model.alias for model in session.agent.config.models]
|
||||
if params.modelId not in model_aliases:
|
||||
return None
|
||||
|
||||
VibeConfig.save_updates({"active_model": params.modelId})
|
||||
|
||||
new_config = VibeConfig.load(
|
||||
workdir=session.agent.config.workdir,
|
||||
tool_paths=session.agent.config.tool_paths,
|
||||
disabled_tools=self._get_disabled_tools_from_capabilities(),
|
||||
)
|
||||
|
||||
await session.agent.reload_with_initial_messages(config=new_config)
|
||||
|
||||
return SetSessionModelResponse()
|
||||
|
||||
@override
|
||||
async def prompt(self, params: PromptRequest) -> PromptResponse:
|
||||
session = self._get_session(params.sessionId)
|
||||
|
||||
if session.task is not None:
|
||||
raise RuntimeError(
|
||||
"Concurrent prompts are not supported yet, wait for agent to finish"
|
||||
)
|
||||
|
||||
text_prompt = self._build_text_prompt(params.prompt)
|
||||
|
||||
async def agent_task() -> None:
|
||||
async for update in self._run_agent_loop(session, text_prompt):
|
||||
await self.connection.sessionUpdate(
|
||||
SessionNotification(sessionId=session.id, update=update)
|
||||
)
|
||||
|
||||
try:
|
||||
session.task = asyncio.create_task(agent_task())
|
||||
await session.task
|
||||
|
||||
except asyncio.CancelledError:
|
||||
return PromptResponse(stopReason="cancelled")
|
||||
|
||||
except Exception as e:
|
||||
await self.connection.sessionUpdate(
|
||||
SessionNotification(
|
||||
sessionId=params.sessionId,
|
||||
update=AgentMessageChunk(
|
||||
sessionUpdate="agent_message_chunk",
|
||||
content=TextContentBlock(type="text", text=f"Error: {e!s}"),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return PromptResponse(stopReason="refusal")
|
||||
|
||||
finally:
|
||||
session.task = None
|
||||
|
||||
return PromptResponse(stopReason="end_turn")
|
||||
|
||||
def _build_text_prompt(self, acp_prompt: list[ContentBlock]) -> str:
|
||||
text_prompt = ""
|
||||
for block in acp_prompt:
|
||||
separator = "\n\n" if text_prompt else ""
|
||||
match block.type:
|
||||
# NOTE: ACP supports annotations, but we don't use them here yet.
|
||||
case "text":
|
||||
text_prompt = f"{text_prompt}{separator}{block.text}"
|
||||
case "resource":
|
||||
block_content = (
|
||||
block.resource.text
|
||||
if isinstance(block.resource, TextResourceContents)
|
||||
else block.resource.blob
|
||||
)
|
||||
fields = {"path": block.resource.uri, "content": block_content}
|
||||
parts = [
|
||||
f"{k}: {v}"
|
||||
for k, v in fields.items()
|
||||
if v is not None and (v or isinstance(v, (int, float)))
|
||||
]
|
||||
block_prompt = "\n".join(parts)
|
||||
text_prompt = f"{text_prompt}{separator}{block_prompt}"
|
||||
case "resource_link":
|
||||
# NOTE: we currently keep more information than just the URI
|
||||
# making it more detailed than the output of the read_file tool.
|
||||
# This is OK, but might be worth testing how it affect performance.
|
||||
fields = {
|
||||
"uri": block.uri,
|
||||
"name": block.name,
|
||||
"title": block.title,
|
||||
"description": block.description,
|
||||
"mimeType": block.mimeType,
|
||||
"size": block.size,
|
||||
}
|
||||
parts = [
|
||||
f"{k}: {v}"
|
||||
for k, v in fields.items()
|
||||
if v is not None and (v or isinstance(v, (int, float)))
|
||||
]
|
||||
block_prompt = "\n".join(parts)
|
||||
text_prompt = f"{text_prompt}{separator}{block_prompt}"
|
||||
case _:
|
||||
raise ValueError(f"Unsupported content block type: {block.type}")
|
||||
return text_prompt
|
||||
|
||||
async def _run_agent_loop(
|
||||
self, session: AcpSession, prompt: str
|
||||
) -> AsyncGenerator[SessionUpdate]:
|
||||
rendered_prompt = render_path_prompt(
|
||||
prompt, base_dir=session.agent.config.effective_workdir
|
||||
)
|
||||
async for event in session.agent.act(rendered_prompt):
|
||||
if isinstance(event, AssistantEvent):
|
||||
yield AgentMessageChunk(
|
||||
sessionUpdate="agent_message_chunk",
|
||||
content=TextContentBlock(type="text", text=event.content),
|
||||
)
|
||||
|
||||
elif isinstance(event, ToolCallEvent):
|
||||
if issubclass(event.tool_class, BaseAcpTool):
|
||||
event.tool_class.update_tool_state(
|
||||
tool_manager=session.agent.tool_manager,
|
||||
connection=self.connection,
|
||||
session_id=session.id,
|
||||
tool_call_id=event.tool_call_id,
|
||||
)
|
||||
|
||||
session_update = tool_call_session_update(event)
|
||||
if session_update:
|
||||
yield session_update
|
||||
|
||||
elif isinstance(event, ToolResultEvent):
|
||||
session_update = tool_result_session_update(event)
|
||||
if session_update:
|
||||
yield session_update
|
||||
|
||||
@override
|
||||
async def cancel(self, params: CancelNotification) -> None:
|
||||
session = self._get_session(params.sessionId)
|
||||
if session.task and not session.task.done():
|
||||
session.task.cancel()
|
||||
session.task = None
|
||||
|
||||
@override
|
||||
async def extMethod(self, method: str, params: dict) -> dict:
|
||||
raise NotImplementedError()
|
||||
|
||||
@override
|
||||
async def extNotification(self, method: str, params: dict) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
async def _run_acp_server() -> None:
|
||||
reader, writer = await stdio_streams()
|
||||
|
||||
AgentSideConnection(lambda connection: VibeAcpAgent(connection), writer, reader)
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
def run_acp_server() -> None:
|
||||
try:
|
||||
asyncio.run(_run_acp_server())
|
||||
except KeyboardInterrupt:
|
||||
# This is expected when the server is terminated
|
||||
pass
|
||||
except Exception as e:
|
||||
# Log any unexpected errors
|
||||
print(f"ACP Agent Server error: {e}", file=sys.stderr)
|
||||
raise
|
||||
37
vibe/acp/entrypoint.py
Normal file
37
vibe/acp/entrypoint.py
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from dataclasses import dataclass
|
||||
import sys
|
||||
|
||||
from vibe.acp.acp_agent import run_acp_server
|
||||
from vibe.setup.onboarding import run_onboarding
|
||||
|
||||
# Configure line buffering for subprocess communication
|
||||
sys.stdout.reconfigure(line_buffering=True) # pyright: ignore[reportAttributeAccessIssue]
|
||||
sys.stderr.reconfigure(line_buffering=True) # pyright: ignore[reportAttributeAccessIssue]
|
||||
sys.stdin.reconfigure(line_buffering=True) # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Arguments:
|
||||
setup: bool
|
||||
|
||||
|
||||
def parse_arguments() -> Arguments:
|
||||
parser = argparse.ArgumentParser(description="Run Mistral Vibe in ACP mode")
|
||||
parser.add_argument("--setup", action="store_true", help="Setup API key and exit")
|
||||
args = parser.parse_args()
|
||||
return Arguments(setup=args.setup)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_arguments()
|
||||
if args.setup:
|
||||
run_onboarding()
|
||||
sys.exit(0)
|
||||
run_acp_server()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
0
vibe/acp/tools/__init__.py
Normal file
0
vibe/acp/tools/__init__.py
Normal file
100
vibe/acp/tools/base.py
Normal file
100
vibe/acp/tools/base.py
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Protocol, cast, runtime_checkable
|
||||
|
||||
from acp import AgentSideConnection, SessionNotification
|
||||
from acp.helpers import SessionUpdate, ToolCallContentVariant
|
||||
from acp.schema import ToolCallProgress
|
||||
from pydantic import Field
|
||||
|
||||
from vibe.core.tools.base import BaseTool, ToolError
|
||||
from vibe.core.tools.manager import ToolManager
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
from vibe.core.utils import logger
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ToolCallSessionUpdateProtocol(Protocol):
|
||||
@classmethod
|
||||
def tool_call_session_update(cls, event: ToolCallEvent) -> SessionUpdate | None: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ToolResultSessionUpdateProtocol(Protocol):
|
||||
@classmethod
|
||||
def tool_result_session_update(
|
||||
cls, event: ToolResultEvent
|
||||
) -> SessionUpdate | None: ...
|
||||
|
||||
|
||||
class AcpToolState:
|
||||
connection: AgentSideConnection | None = Field(
|
||||
default=None, description="ACP agent-side connection"
|
||||
)
|
||||
session_id: str | None = Field(default=None, description="Current ACP session ID")
|
||||
tool_call_id: str | None = Field(
|
||||
default=None, description="Current ACP tool call ID"
|
||||
)
|
||||
|
||||
|
||||
class BaseAcpTool[ToolState: AcpToolState](BaseTool):
|
||||
state: ToolState
|
||||
|
||||
@classmethod
|
||||
def get_tool_instance(
|
||||
cls, tool_name: str, tool_manager: ToolManager
|
||||
) -> BaseAcpTool[AcpToolState]:
|
||||
return cast(BaseAcpTool[AcpToolState], tool_manager.get(tool_name))
|
||||
|
||||
@classmethod
|
||||
def update_tool_state(
|
||||
cls,
|
||||
*,
|
||||
tool_manager: ToolManager,
|
||||
connection: AgentSideConnection | None,
|
||||
session_id: str | None,
|
||||
tool_call_id: str | None,
|
||||
) -> None:
|
||||
tool_instance = cls.get_tool_instance(cls.get_name(), tool_manager)
|
||||
tool_instance.state.connection = connection
|
||||
tool_instance.state.session_id = session_id
|
||||
tool_instance.state.tool_call_id = tool_call_id
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def _get_tool_state_class(cls) -> type[ToolState]: ...
|
||||
|
||||
def _load_state(self) -> tuple[AgentSideConnection, str, str | None]:
|
||||
if self.state.connection is None:
|
||||
raise ToolError(
|
||||
"Connection not available in tool state. This tool can only be used within an ACP session."
|
||||
)
|
||||
if self.state.session_id is None:
|
||||
raise ToolError(
|
||||
"Session ID not available in tool state. This tool can only be used within an ACP session."
|
||||
)
|
||||
|
||||
return self.state.connection, self.state.session_id, self.state.tool_call_id
|
||||
|
||||
async def _send_in_progress_session_update(
|
||||
self, content: list[ToolCallContentVariant] | None = None
|
||||
) -> None:
|
||||
connection, session_id, tool_call_id = self._load_state()
|
||||
if tool_call_id is None:
|
||||
return
|
||||
|
||||
try:
|
||||
await connection.sessionUpdate(
|
||||
SessionNotification(
|
||||
sessionId=session_id,
|
||||
update=ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=tool_call_id,
|
||||
status="in_progress",
|
||||
content=content,
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update session: {e!r}")
|
||||
144
vibe/acp/tools/builtins/bash.py
Normal file
144
vibe/acp/tools/builtins/bash.py
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shlex
|
||||
|
||||
from acp import CreateTerminalRequest, TerminalHandle
|
||||
from acp.schema import (
|
||||
EnvVariable,
|
||||
TerminalToolCallContent,
|
||||
ToolCallProgress,
|
||||
ToolCallStart,
|
||||
WaitForTerminalExitResponse,
|
||||
)
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import AcpToolState, BaseAcpTool
|
||||
from vibe.core.tools.base import BaseToolState, ToolError
|
||||
from vibe.core.tools.builtins.bash import Bash as CoreBashTool, BashArgs, BashResult
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
from vibe.core.utils import logger
|
||||
|
||||
|
||||
class AcpBashState(BaseToolState, AcpToolState):
|
||||
pass
|
||||
|
||||
|
||||
class Bash(CoreBashTool, BaseAcpTool[AcpBashState]):
|
||||
prompt_path = VIBE_ROOT / "core" / "tools" / "builtins" / "prompts" / "bash.md"
|
||||
state: AcpBashState
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[AcpBashState]:
|
||||
return AcpBashState
|
||||
|
||||
async def run(self, args: BashArgs) -> BashResult:
|
||||
connection, session_id, _ = self._load_state()
|
||||
|
||||
timeout = args.timeout or self.config.default_timeout
|
||||
max_bytes = self.config.max_output_bytes
|
||||
env, command, cmd_args = self._parse_command(args.command)
|
||||
|
||||
create_request = CreateTerminalRequest(
|
||||
sessionId=session_id,
|
||||
command=command,
|
||||
args=cmd_args,
|
||||
env=env,
|
||||
cwd=str(self.config.effective_workdir),
|
||||
outputByteLimit=max_bytes,
|
||||
)
|
||||
|
||||
try:
|
||||
terminal_handle = await connection.createTerminal(create_request)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Failed to create terminal: {e!r}") from e
|
||||
|
||||
await self._send_in_progress_session_update([
|
||||
TerminalToolCallContent(type="terminal", terminalId=terminal_handle.id)
|
||||
])
|
||||
|
||||
try:
|
||||
exit_response = await self._wait_for_terminal_exit(
|
||||
terminal_handle, timeout, args.command
|
||||
)
|
||||
|
||||
output_response = await terminal_handle.current_output()
|
||||
|
||||
return self._build_result(
|
||||
command=args.command,
|
||||
stdout=output_response.output,
|
||||
stderr="",
|
||||
returncode=exit_response.exitCode or 0,
|
||||
)
|
||||
|
||||
finally:
|
||||
try:
|
||||
await terminal_handle.release()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to release terminal: {e!r}")
|
||||
|
||||
def _parse_command(
|
||||
self, command_str: str
|
||||
) -> tuple[list[EnvVariable], str, list[str]]:
|
||||
parts = shlex.split(command_str)
|
||||
env: list[EnvVariable] = []
|
||||
command: str = ""
|
||||
args: list[str] = []
|
||||
|
||||
for part in parts:
|
||||
if "=" in part and not command:
|
||||
key, value = part.split("=", 1)
|
||||
env.append(EnvVariable(name=key, value=value))
|
||||
elif not command:
|
||||
command = part
|
||||
else:
|
||||
args.append(part)
|
||||
|
||||
return env, command, args
|
||||
|
||||
@classmethod
|
||||
def get_summary(cls, args: BashArgs) -> str:
|
||||
summary = f"{args.command}"
|
||||
if args.timeout:
|
||||
summary += f" (timeout {args.timeout}s)"
|
||||
|
||||
return summary
|
||||
|
||||
async def _wait_for_terminal_exit(
|
||||
self, terminal_handle: TerminalHandle, timeout: int, command: str
|
||||
) -> WaitForTerminalExitResponse:
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
terminal_handle.wait_for_exit(), timeout=timeout
|
||||
)
|
||||
except TimeoutError:
|
||||
try:
|
||||
await terminal_handle.kill()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to kill terminal: {e!r}")
|
||||
|
||||
raise self._build_timeout_error(command, timeout)
|
||||
|
||||
@classmethod
|
||||
def tool_call_session_update(cls, event: ToolCallEvent) -> ToolCallStart:
|
||||
if not isinstance(event.args, BashArgs):
|
||||
raise ValueError(f"Unexpected tool args: {event.args}")
|
||||
|
||||
return ToolCallStart(
|
||||
sessionUpdate="tool_call",
|
||||
title=Bash.get_summary(event.args),
|
||||
content=None,
|
||||
toolCallId=event.tool_call_id,
|
||||
kind="execute",
|
||||
rawInput=event.args.model_dump_json(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tool_result_session_update(
|
||||
cls, event: ToolResultEvent
|
||||
) -> ToolCallProgress | None:
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="failed" if event.error else "completed",
|
||||
)
|
||||
58
vibe/acp/tools/builtins/read_file.py
Normal file
58
vibe/acp/tools/builtins/read_file.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from acp import ReadTextFileRequest
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import AcpToolState, BaseAcpTool
|
||||
from vibe.core.tools.base import ToolError
|
||||
from vibe.core.tools.builtins.read_file import (
|
||||
ReadFile as CoreReadFileTool,
|
||||
ReadFileArgs,
|
||||
ReadFileResult,
|
||||
ReadFileState,
|
||||
_ReadResult,
|
||||
)
|
||||
|
||||
ReadFileResult = ReadFileResult
|
||||
|
||||
|
||||
class AcpReadFileState(ReadFileState, AcpToolState):
|
||||
pass
|
||||
|
||||
|
||||
class ReadFile(CoreReadFileTool, BaseAcpTool[AcpReadFileState]):
|
||||
state: AcpReadFileState
|
||||
prompt_path = VIBE_ROOT / "core" / "tools" / "builtins" / "prompts" / "read_file.md"
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[AcpReadFileState]:
|
||||
return AcpReadFileState
|
||||
|
||||
async def _read_file(self, args: ReadFileArgs, file_path: Path) -> _ReadResult:
|
||||
connection, session_id, _ = self._load_state()
|
||||
|
||||
line = args.offset + 1 if args.offset > 0 else None
|
||||
limit = args.limit
|
||||
|
||||
read_request = ReadTextFileRequest(
|
||||
sessionId=session_id, path=str(file_path), line=line, limit=limit
|
||||
)
|
||||
|
||||
await self._send_in_progress_session_update()
|
||||
|
||||
try:
|
||||
response = await connection.readTextFile(read_request)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error reading {file_path}: {e}") from e
|
||||
|
||||
content_lines = response.content.splitlines(keepends=True)
|
||||
lines_read = len(content_lines)
|
||||
bytes_read = sum(len(line.encode("utf-8")) for line in content_lines)
|
||||
|
||||
was_truncated = args.limit is not None and lines_read >= args.limit
|
||||
|
||||
return _ReadResult(
|
||||
lines=content_lines, bytes_read=bytes_read, was_truncated=was_truncated
|
||||
)
|
||||
132
vibe/acp/tools/builtins/search_replace.py
Normal file
132
vibe/acp/tools/builtins/search_replace.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from acp import ReadTextFileRequest, WriteTextFileRequest
|
||||
from acp.helpers import SessionUpdate
|
||||
from acp.schema import (
|
||||
FileEditToolCallContent,
|
||||
ToolCallLocation,
|
||||
ToolCallProgress,
|
||||
ToolCallStart,
|
||||
)
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import AcpToolState, BaseAcpTool
|
||||
from vibe.core.tools.base import ToolError
|
||||
from vibe.core.tools.builtins.search_replace import (
|
||||
SearchReplace as CoreSearchReplaceTool,
|
||||
SearchReplaceArgs,
|
||||
SearchReplaceResult,
|
||||
SearchReplaceState,
|
||||
)
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class AcpSearchReplaceState(SearchReplaceState, AcpToolState):
|
||||
file_backup_content: str | None = None
|
||||
|
||||
|
||||
class SearchReplace(CoreSearchReplaceTool, BaseAcpTool[AcpSearchReplaceState]):
|
||||
state: AcpSearchReplaceState
|
||||
prompt_path = (
|
||||
VIBE_ROOT / "core" / "tools" / "builtins" / "prompts" / "search_replace.md"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[AcpSearchReplaceState]:
|
||||
return AcpSearchReplaceState
|
||||
|
||||
async def _read_file(self, file_path: Path) -> str:
|
||||
connection, session_id, _ = self._load_state()
|
||||
|
||||
read_request = ReadTextFileRequest(sessionId=session_id, path=str(file_path))
|
||||
|
||||
await self._send_in_progress_session_update()
|
||||
|
||||
try:
|
||||
response = await connection.readTextFile(read_request)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Unexpected error reading {file_path}: {e}") from e
|
||||
|
||||
self.state.file_backup_content = response.content
|
||||
return response.content
|
||||
|
||||
async def _backup_file(self, file_path: Path) -> None:
|
||||
if self.state.file_backup_content is None:
|
||||
return
|
||||
|
||||
await self._write_file(
|
||||
file_path.with_suffix(file_path.suffix + ".bak"),
|
||||
self.state.file_backup_content,
|
||||
)
|
||||
|
||||
async def _write_file(self, file_path: Path, content: str) -> None:
|
||||
connection, session_id, _ = self._load_state()
|
||||
|
||||
write_request = WriteTextFileRequest(
|
||||
sessionId=session_id, path=str(file_path), content=content
|
||||
)
|
||||
|
||||
try:
|
||||
await connection.writeTextFile(write_request)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error writing {file_path}: {e}") from e
|
||||
|
||||
@classmethod
|
||||
def tool_call_session_update(cls, event: ToolCallEvent) -> SessionUpdate | None:
|
||||
args = event.args
|
||||
if not isinstance(args, SearchReplaceArgs):
|
||||
return None
|
||||
|
||||
blocks = cls._parse_search_replace_blocks(args.content)
|
||||
|
||||
return ToolCallStart(
|
||||
sessionUpdate="tool_call",
|
||||
title=cls.get_call_display(event).summary,
|
||||
toolCallId=event.tool_call_id,
|
||||
kind="edit",
|
||||
content=[
|
||||
FileEditToolCallContent(
|
||||
type="diff",
|
||||
path=args.file_path,
|
||||
oldText=block.search,
|
||||
newText=block.replace,
|
||||
)
|
||||
for block in blocks
|
||||
],
|
||||
locations=[ToolCallLocation(path=args.file_path)],
|
||||
rawInput=args.model_dump_json(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tool_result_session_update(cls, event: ToolResultEvent) -> SessionUpdate | None:
|
||||
if event.error:
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="failed",
|
||||
)
|
||||
|
||||
result = event.result
|
||||
if not isinstance(result, SearchReplaceResult):
|
||||
return None
|
||||
|
||||
blocks = cls._parse_search_replace_blocks(result.content)
|
||||
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="completed",
|
||||
content=[
|
||||
FileEditToolCallContent(
|
||||
type="diff",
|
||||
path=result.file,
|
||||
oldText=block.search,
|
||||
newText=block.replace,
|
||||
)
|
||||
for block in blocks
|
||||
],
|
||||
locations=[ToolCallLocation(path=result.file)],
|
||||
rawOutput=result.model_dump_json(),
|
||||
)
|
||||
65
vibe/acp/tools/builtins/todo.py
Normal file
65
vibe/acp/tools/builtins/todo.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import cast
|
||||
|
||||
from acp.helpers import SessionUpdate
|
||||
from acp.schema import AgentPlanUpdate, PlanEntry, PlanEntryPriority, PlanEntryStatus
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import AcpToolState, BaseAcpTool
|
||||
from vibe.core.tools.builtins.todo import (
|
||||
Todo as CoreTodoTool,
|
||||
TodoArgs,
|
||||
TodoPriority,
|
||||
TodoResult,
|
||||
TodoState,
|
||||
TodoStatus,
|
||||
)
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
TodoArgs = TodoArgs
|
||||
|
||||
|
||||
class AcpTodoState(TodoState, AcpToolState):
|
||||
pass
|
||||
|
||||
|
||||
class Todo(CoreTodoTool, BaseAcpTool[AcpTodoState]):
|
||||
state: AcpTodoState
|
||||
prompt_path = VIBE_ROOT / "core" / "tools" / "builtins" / "prompts" / "todo.md"
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[AcpTodoState]:
|
||||
return AcpTodoState
|
||||
|
||||
@classmethod
|
||||
def tool_call_session_update(cls, event: ToolCallEvent) -> SessionUpdate | None:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def tool_result_session_update(cls, event: ToolResultEvent) -> SessionUpdate | None:
|
||||
result = cast(TodoResult, event.result)
|
||||
todos = [todo for todo in result.todos if todo.status != TodoStatus.CANCELLED]
|
||||
matched_status: dict[TodoStatus, PlanEntryStatus] = {
|
||||
TodoStatus.PENDING: "pending",
|
||||
TodoStatus.IN_PROGRESS: "in_progress",
|
||||
TodoStatus.COMPLETED: "completed",
|
||||
}
|
||||
matched_priority: dict[TodoPriority, PlanEntryPriority] = {
|
||||
TodoPriority.LOW: "low",
|
||||
TodoPriority.MEDIUM: "medium",
|
||||
TodoPriority.HIGH: "high",
|
||||
}
|
||||
|
||||
update = AgentPlanUpdate(
|
||||
sessionUpdate="plan",
|
||||
entries=[
|
||||
PlanEntry(
|
||||
content=todo.content,
|
||||
status=matched_status[todo.status],
|
||||
priority=matched_priority[todo.priority],
|
||||
)
|
||||
for todo in todos
|
||||
],
|
||||
)
|
||||
return update
|
||||
98
vibe/acp/tools/builtins/write_file.py
Normal file
98
vibe/acp/tools/builtins/write_file.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from acp import WriteTextFileRequest
|
||||
from acp.helpers import SessionUpdate
|
||||
from acp.schema import (
|
||||
FileEditToolCallContent,
|
||||
ToolCallLocation,
|
||||
ToolCallProgress,
|
||||
ToolCallStart,
|
||||
)
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.acp.tools.base import AcpToolState, BaseAcpTool
|
||||
from vibe.core.tools.base import ToolError
|
||||
from vibe.core.tools.builtins.write_file import (
|
||||
WriteFile as CoreWriteFileTool,
|
||||
WriteFileArgs,
|
||||
WriteFileResult,
|
||||
WriteFileState,
|
||||
)
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class AcpWriteFileState(WriteFileState, AcpToolState):
|
||||
pass
|
||||
|
||||
|
||||
class WriteFile(CoreWriteFileTool, BaseAcpTool[AcpWriteFileState]):
|
||||
state: AcpWriteFileState
|
||||
prompt_path = (
|
||||
VIBE_ROOT / "core" / "tools" / "builtins" / "prompts" / "write_file.md"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[AcpWriteFileState]:
|
||||
return AcpWriteFileState
|
||||
|
||||
async def _write_file(self, args: WriteFileArgs, file_path: Path) -> None:
|
||||
connection, session_id, _ = self._load_state()
|
||||
|
||||
write_request = WriteTextFileRequest(
|
||||
sessionId=session_id, path=str(file_path), content=args.content
|
||||
)
|
||||
|
||||
await self._send_in_progress_session_update()
|
||||
|
||||
try:
|
||||
await connection.writeTextFile(write_request)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error writing {file_path}: {e}") from e
|
||||
|
||||
@classmethod
|
||||
def tool_call_session_update(cls, event: ToolCallEvent) -> SessionUpdate | None:
|
||||
args = event.args
|
||||
if not isinstance(args, WriteFileArgs):
|
||||
return None
|
||||
|
||||
return ToolCallStart(
|
||||
sessionUpdate="tool_call",
|
||||
title=cls.get_call_display(event).summary,
|
||||
toolCallId=event.tool_call_id,
|
||||
kind="edit",
|
||||
content=[
|
||||
FileEditToolCallContent(
|
||||
type="diff", path=args.path, oldText=None, newText=args.content
|
||||
)
|
||||
],
|
||||
locations=[ToolCallLocation(path=args.path)],
|
||||
rawInput=args.model_dump_json(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tool_result_session_update(cls, event: ToolResultEvent) -> SessionUpdate | None:
|
||||
if event.error:
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="failed",
|
||||
)
|
||||
|
||||
result = event.result
|
||||
if not isinstance(result, WriteFileResult):
|
||||
return None
|
||||
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="completed",
|
||||
content=[
|
||||
FileEditToolCallContent(
|
||||
type="diff", path=result.path, oldText=None, newText=result.content
|
||||
)
|
||||
],
|
||||
locations=[ToolCallLocation(path=result.path)],
|
||||
rawOutput=result.model_dump_json(),
|
||||
)
|
||||
111
vibe/acp/tools/session_update.py
Normal file
111
vibe/acp/tools/session_update.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from acp.helpers import SessionUpdate, ToolCallContentVariant
|
||||
from acp.schema import (
|
||||
ContentToolCallContent,
|
||||
TextContentBlock,
|
||||
ToolCallProgress,
|
||||
ToolCallStart,
|
||||
ToolKind,
|
||||
)
|
||||
|
||||
from vibe.acp.tools.base import (
|
||||
ToolCallSessionUpdateProtocol,
|
||||
ToolResultSessionUpdateProtocol,
|
||||
)
|
||||
from vibe.core.tools.ui import ToolUIDataAdapter
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
from vibe.core.utils import TaggedText, is_user_cancellation_event
|
||||
|
||||
TOOL_KIND: dict[str, ToolKind] = {"read_file": "read", "grep": "search"}
|
||||
|
||||
|
||||
def tool_call_session_update(event: ToolCallEvent) -> SessionUpdate | None:
|
||||
if issubclass(event.tool_class, ToolCallSessionUpdateProtocol):
|
||||
return event.tool_class.tool_call_session_update(event)
|
||||
|
||||
adapter = ToolUIDataAdapter(event.tool_class)
|
||||
display = adapter.get_call_display(event)
|
||||
content: list[ToolCallContentVariant] | None = (
|
||||
[
|
||||
ContentToolCallContent(
|
||||
type="content",
|
||||
content=TextContentBlock(type="text", text=display.content),
|
||||
)
|
||||
]
|
||||
if display.content
|
||||
else None
|
||||
)
|
||||
|
||||
return ToolCallStart(
|
||||
sessionUpdate="tool_call",
|
||||
title=display.summary,
|
||||
content=content,
|
||||
toolCallId=event.tool_call_id,
|
||||
kind=TOOL_KIND.get(event.tool_name, "other"),
|
||||
rawInput=event.args.model_dump_json(),
|
||||
)
|
||||
|
||||
|
||||
def tool_result_session_update(event: ToolResultEvent) -> SessionUpdate | None:
|
||||
if is_user_cancellation_event(event):
|
||||
tool_status = "failed"
|
||||
raw_output = (
|
||||
TaggedText.from_string(event.skip_reason).message
|
||||
if event.skip_reason
|
||||
else None
|
||||
)
|
||||
elif event.result:
|
||||
tool_status = "completed"
|
||||
raw_output = event.result.model_dump_json()
|
||||
else:
|
||||
tool_status = "failed"
|
||||
raw_output = (
|
||||
TaggedText.from_string(event.error).message if event.error else None
|
||||
)
|
||||
|
||||
if event.tool_class is None:
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status="failed",
|
||||
rawOutput=raw_output,
|
||||
content=[
|
||||
ContentToolCallContent(
|
||||
type="content",
|
||||
content=TextContentBlock(type="text", text=raw_output or ""),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
if issubclass(event.tool_class, ToolResultSessionUpdateProtocol):
|
||||
return event.tool_class.tool_result_session_update(event)
|
||||
|
||||
if tool_status == "failed":
|
||||
content = [
|
||||
ContentToolCallContent(
|
||||
type="content",
|
||||
content=TextContentBlock(type="text", text=raw_output or ""),
|
||||
)
|
||||
]
|
||||
else:
|
||||
adapter = ToolUIDataAdapter(event.tool_class)
|
||||
display = adapter.get_result_display(event)
|
||||
content: list[ToolCallContentVariant] | None = (
|
||||
[
|
||||
ContentToolCallContent(
|
||||
type="content",
|
||||
content=TextContentBlock(type="text", text=display.message),
|
||||
)
|
||||
]
|
||||
if display.message
|
||||
else None
|
||||
)
|
||||
|
||||
return ToolCallProgress(
|
||||
sessionUpdate="tool_call_update",
|
||||
toolCallId=event.tool_call_id,
|
||||
status=tool_status,
|
||||
rawOutput=raw_output,
|
||||
content=content,
|
||||
)
|
||||
70
vibe/acp/utils.py
Normal file
70
vibe/acp/utils.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
from enum import StrEnum
|
||||
from typing import Literal, cast
|
||||
|
||||
from acp.schema import PermissionOption, SessionMode
|
||||
|
||||
|
||||
class VibeSessionMode(enum.StrEnum):
|
||||
APPROVAL_REQUIRED = enum.auto()
|
||||
AUTO_APPROVE = enum.auto()
|
||||
|
||||
def to_acp_session_mode(self) -> SessionMode:
|
||||
match self:
|
||||
case self.APPROVAL_REQUIRED:
|
||||
return SessionMode(
|
||||
id=VibeSessionMode.APPROVAL_REQUIRED,
|
||||
name="Approval Required",
|
||||
description="Requires user approval for tool executions",
|
||||
)
|
||||
case self.AUTO_APPROVE:
|
||||
return SessionMode(
|
||||
id=VibeSessionMode.AUTO_APPROVE,
|
||||
name="Auto Approve",
|
||||
description="Automatically approves all tool executions",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_acp_session_mode(cls, session_mode: SessionMode) -> VibeSessionMode | None:
|
||||
if not cls.is_valid(session_mode.id):
|
||||
return None
|
||||
return cls(session_mode.id)
|
||||
|
||||
@classmethod
|
||||
def is_valid(cls, mode_id: str) -> bool:
|
||||
try:
|
||||
return cls(mode_id).to_acp_session_mode() is not None
|
||||
except (ValueError, KeyError):
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_all_acp_session_modes(cls) -> list[SessionMode]:
|
||||
return [mode.to_acp_session_mode() for mode in cls]
|
||||
|
||||
|
||||
class ToolOption(StrEnum):
|
||||
ALLOW_ONCE = "allow_once"
|
||||
ALLOW_ALWAYS = "allow_always"
|
||||
REJECT_ONCE = "reject_once"
|
||||
REJECT_ALWAYS = "reject_always"
|
||||
|
||||
|
||||
TOOL_OPTIONS = [
|
||||
PermissionOption(
|
||||
optionId=ToolOption.ALLOW_ONCE,
|
||||
name="Allow once",
|
||||
kind=cast(Literal["allow_once"], ToolOption.ALLOW_ONCE),
|
||||
),
|
||||
PermissionOption(
|
||||
optionId=ToolOption.ALLOW_ALWAYS,
|
||||
name="Allow always",
|
||||
kind=cast(Literal["allow_always"], ToolOption.ALLOW_ALWAYS),
|
||||
),
|
||||
PermissionOption(
|
||||
optionId=ToolOption.REJECT_ONCE,
|
||||
name="Reject once",
|
||||
kind=cast(Literal["reject_once"], ToolOption.REJECT_ONCE),
|
||||
),
|
||||
]
|
||||
0
vibe/cli/__init__.py
Normal file
0
vibe/cli/__init__.py
Normal file
0
vibe/cli/autocompletion/__init__.py
Normal file
0
vibe/cli/autocompletion/__init__.py
Normal file
22
vibe/cli/autocompletion/base.py
Normal file
22
vibe/cli/autocompletion/base.py
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Protocol
|
||||
|
||||
|
||||
class CompletionResult(StrEnum):
|
||||
IGNORED = "ignored"
|
||||
HANDLED = "handled"
|
||||
SUBMIT = "submit"
|
||||
|
||||
|
||||
class CompletionView(Protocol):
|
||||
def render_completion_suggestions(
|
||||
self, suggestions: list[tuple[str, str]], selected_index: int
|
||||
) -> None: ...
|
||||
|
||||
def clear_completion_suggestions(self) -> None: ...
|
||||
|
||||
def replace_completion_range(
|
||||
self, start: int, end: int, replacement: str
|
||||
) -> None: ...
|
||||
173
vibe/cli/autocompletion/path_completion.py
Normal file
173
vibe/cli/autocompletion/path_completion.py
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from threading import Lock
|
||||
|
||||
from textual import events
|
||||
|
||||
from vibe.cli.autocompletion.base import CompletionResult, CompletionView
|
||||
from vibe.core.autocompletion.completers import PathCompleter
|
||||
|
||||
MAX_SUGGESTIONS_COUNT = 10
|
||||
|
||||
|
||||
class PathCompletionController:
|
||||
_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="path-completion")
|
||||
|
||||
def __init__(self, completer: PathCompleter, view: CompletionView) -> None:
|
||||
self._completer = completer
|
||||
self._view = view
|
||||
self._suggestions: list[tuple[str, str]] = []
|
||||
self._selected_index = 0
|
||||
self._pending_future: Future | None = None
|
||||
self._last_query: tuple[str, int] | None = None
|
||||
self._query_lock = Lock()
|
||||
|
||||
def can_handle(self, text: str, cursor_index: int) -> bool:
|
||||
if cursor_index < 0 or cursor_index > len(text):
|
||||
return False
|
||||
|
||||
if cursor_index == 0:
|
||||
return False
|
||||
|
||||
before_cursor = text[:cursor_index]
|
||||
if "@" not in before_cursor:
|
||||
return False
|
||||
|
||||
at_index = before_cursor.rfind("@")
|
||||
|
||||
if cursor_index <= at_index:
|
||||
return False
|
||||
|
||||
fragment = before_cursor[at_index:cursor_index]
|
||||
# fragment must not be empty (including @) and not contain any spaces
|
||||
return bool(fragment) and " " not in fragment
|
||||
|
||||
def reset(self) -> None:
|
||||
with self._query_lock:
|
||||
if self._pending_future and not self._pending_future.done():
|
||||
self._pending_future.cancel()
|
||||
self._pending_future = None
|
||||
self._last_query = None
|
||||
if self._suggestions:
|
||||
self._suggestions.clear()
|
||||
self._selected_index = 0
|
||||
self._view.clear_completion_suggestions()
|
||||
|
||||
def on_text_changed(self, text: str, cursor_index: int) -> None:
|
||||
if not self.can_handle(text, cursor_index):
|
||||
self.reset()
|
||||
return
|
||||
|
||||
query = (text, cursor_index)
|
||||
with self._query_lock:
|
||||
if query == self._last_query:
|
||||
return
|
||||
|
||||
if self._pending_future and not self._pending_future.done():
|
||||
# NOTE (Vince): this is a "best effort" cancellation: it only works if the task
|
||||
# hasn't started; once running in the thread pool, it cannot be cancelled
|
||||
self._pending_future.cancel()
|
||||
|
||||
self._last_query = query
|
||||
|
||||
app = getattr(self._view, "app", None)
|
||||
if app:
|
||||
with self._query_lock:
|
||||
self._pending_future = self._executor.submit(
|
||||
self._compute_completions, text, cursor_index
|
||||
)
|
||||
self._pending_future.add_done_callback(
|
||||
lambda f: self._handle_completion_result(f, query)
|
||||
)
|
||||
else:
|
||||
suggestions = self._compute_completions(text, cursor_index)
|
||||
self._update_suggestions(suggestions)
|
||||
|
||||
def _compute_completions(
|
||||
self, text: str, cursor_index: int
|
||||
) -> list[tuple[str, str]]:
|
||||
return self._completer.get_completion_items(text, cursor_index)
|
||||
|
||||
def _handle_completion_result(self, future: Future, query: tuple[str, int]) -> None:
|
||||
if future.cancelled():
|
||||
return
|
||||
|
||||
try:
|
||||
suggestions = future.result()
|
||||
with self._query_lock:
|
||||
if query == self._last_query:
|
||||
self._update_suggestions(suggestions)
|
||||
except Exception:
|
||||
with self._query_lock:
|
||||
self._pending_future = None
|
||||
self._last_query = None
|
||||
|
||||
def _update_suggestions(self, suggestions: list[tuple[str, str]]) -> None:
|
||||
if len(suggestions) > MAX_SUGGESTIONS_COUNT:
|
||||
suggestions = suggestions[:MAX_SUGGESTIONS_COUNT]
|
||||
|
||||
app = getattr(self._view, "app", None)
|
||||
|
||||
if suggestions:
|
||||
self._suggestions = suggestions
|
||||
self._selected_index = 0
|
||||
if app:
|
||||
app.call_after_refresh(
|
||||
self._view.render_completion_suggestions,
|
||||
self._suggestions,
|
||||
self._selected_index,
|
||||
)
|
||||
else:
|
||||
self._view.render_completion_suggestions(
|
||||
self._suggestions, self._selected_index
|
||||
)
|
||||
elif app:
|
||||
app.call_after_refresh(self.reset)
|
||||
else:
|
||||
self.reset()
|
||||
|
||||
def on_key(
|
||||
self, event: events.Key, text: str, cursor_index: int
|
||||
) -> CompletionResult:
|
||||
if not self._suggestions:
|
||||
return CompletionResult.IGNORED
|
||||
|
||||
match event.key:
|
||||
case "tab" | "enter":
|
||||
if self._apply_selected_completion(text, cursor_index):
|
||||
return CompletionResult.HANDLED
|
||||
return CompletionResult.IGNORED
|
||||
case "down":
|
||||
self._move_selection(1)
|
||||
return CompletionResult.HANDLED
|
||||
case "up":
|
||||
self._move_selection(-1)
|
||||
return CompletionResult.HANDLED
|
||||
case _:
|
||||
return CompletionResult.IGNORED
|
||||
|
||||
def _move_selection(self, delta: int) -> None:
|
||||
if not self._suggestions:
|
||||
return
|
||||
|
||||
count = len(self._suggestions)
|
||||
self._selected_index = (self._selected_index + delta) % count
|
||||
self._view.render_completion_suggestions(
|
||||
self._suggestions, self._selected_index
|
||||
)
|
||||
|
||||
def _apply_selected_completion(self, text: str, cursor_index: int) -> bool:
|
||||
if not self._suggestions:
|
||||
return False
|
||||
|
||||
completion, _ = self._suggestions[self._selected_index]
|
||||
replacement_range = self._completer.get_replacement_range(text, cursor_index)
|
||||
if replacement_range is None:
|
||||
self.reset()
|
||||
return False
|
||||
|
||||
start, end = replacement_range
|
||||
self._view.replace_completion_range(start, end, completion)
|
||||
self.reset()
|
||||
return True
|
||||
99
vibe/cli/autocompletion/slash_command.py
Normal file
99
vibe/cli/autocompletion/slash_command.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from textual import events
|
||||
|
||||
from vibe.cli.autocompletion.base import CompletionResult, CompletionView
|
||||
from vibe.core.autocompletion.completers import CommandCompleter
|
||||
|
||||
MAX_SUGGESTIONS_COUNT = 5
|
||||
|
||||
|
||||
class SlashCommandController:
|
||||
def __init__(self, completer: CommandCompleter, view: CompletionView) -> None:
|
||||
self._completer = completer
|
||||
self._view = view
|
||||
self._suggestions: list[tuple[str, str]] = []
|
||||
self._selected_index = 0
|
||||
|
||||
def can_handle(self, text: str, cursor_index: int) -> bool:
|
||||
return text.startswith("/")
|
||||
|
||||
def reset(self) -> None:
|
||||
if self._suggestions:
|
||||
self._suggestions.clear()
|
||||
self._selected_index = 0
|
||||
self._view.clear_completion_suggestions()
|
||||
|
||||
def on_text_changed(self, text: str, cursor_index: int) -> None:
|
||||
if cursor_index < 0 or cursor_index > len(text):
|
||||
self.reset()
|
||||
return
|
||||
|
||||
if not self.can_handle(text, cursor_index):
|
||||
self.reset()
|
||||
return
|
||||
|
||||
suggestions = self._completer.get_completion_items(text, cursor_index)
|
||||
if len(suggestions) > MAX_SUGGESTIONS_COUNT:
|
||||
suggestions = suggestions[:MAX_SUGGESTIONS_COUNT]
|
||||
if suggestions:
|
||||
self._suggestions = suggestions
|
||||
self._selected_index = 0
|
||||
self._view.render_completion_suggestions(
|
||||
self._suggestions, self._selected_index
|
||||
)
|
||||
else:
|
||||
self.reset()
|
||||
|
||||
def on_key(
|
||||
self, event: events.Key, text: str, cursor_index: int
|
||||
) -> CompletionResult:
|
||||
if not self._suggestions:
|
||||
return CompletionResult.IGNORED
|
||||
|
||||
match event.key:
|
||||
case "tab":
|
||||
if self._apply_selected_completion(text, cursor_index):
|
||||
result = CompletionResult.HANDLED
|
||||
else:
|
||||
result = CompletionResult.IGNORED
|
||||
case "enter":
|
||||
if self._apply_selected_completion(text, cursor_index):
|
||||
result = CompletionResult.SUBMIT
|
||||
else:
|
||||
result = CompletionResult.HANDLED
|
||||
case "down":
|
||||
self._move_selection(1)
|
||||
result = CompletionResult.HANDLED
|
||||
case "up":
|
||||
self._move_selection(-1)
|
||||
result = CompletionResult.HANDLED
|
||||
case _:
|
||||
result = CompletionResult.IGNORED
|
||||
|
||||
return result
|
||||
|
||||
def _move_selection(self, delta: int) -> None:
|
||||
if not self._suggestions:
|
||||
return
|
||||
|
||||
count = len(self._suggestions)
|
||||
self._selected_index = (self._selected_index + delta) % count
|
||||
self._view.render_completion_suggestions(
|
||||
self._suggestions, self._selected_index
|
||||
)
|
||||
|
||||
def _apply_selected_completion(self, text: str, cursor_index: int) -> bool:
|
||||
if not self._suggestions:
|
||||
return False
|
||||
|
||||
alias, _ = self._suggestions[self._selected_index]
|
||||
replacement_range = self._completer.get_replacement_range(text, cursor_index)
|
||||
if replacement_range is None:
|
||||
self.reset()
|
||||
return False
|
||||
|
||||
start, end = replacement_range
|
||||
self._view.replace_completion_range(start, end, alias)
|
||||
self.reset()
|
||||
return True
|
||||
34
vibe/cli/clipboard.py
Normal file
34
vibe/cli/clipboard.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pyperclip
|
||||
from textual.app import App
|
||||
|
||||
|
||||
def copy_selection_to_clipboard(app: App) -> None:
|
||||
selected_texts = []
|
||||
|
||||
for widget in app.query("*"):
|
||||
if not hasattr(widget, "text_selection") or not widget.text_selection:
|
||||
continue
|
||||
|
||||
selection = widget.text_selection
|
||||
result = widget.get_selection(selection)
|
||||
if not result:
|
||||
continue
|
||||
|
||||
selected_text, _ = result
|
||||
if selected_text.strip():
|
||||
selected_texts.append(selected_text)
|
||||
|
||||
if not selected_texts:
|
||||
return
|
||||
|
||||
combined_text = "\n".join(selected_texts)
|
||||
|
||||
try:
|
||||
pyperclip.copy(combined_text)
|
||||
app.notify("Selection added to clipboard", severity="information", timeout=2)
|
||||
except Exception:
|
||||
app.notify(
|
||||
"Use Ctrl+c to copy selections in Vibe", severity="warning", timeout=3
|
||||
)
|
||||
98
vibe/cli/commands.py
Normal file
98
vibe/cli/commands.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Command:
|
||||
aliases: frozenset[str]
|
||||
description: str
|
||||
handler: str
|
||||
exits: bool = False
|
||||
|
||||
|
||||
class CommandRegistry:
|
||||
def __init__(self, excluded_commands: list[str] | None = None) -> None:
|
||||
if excluded_commands is None:
|
||||
excluded_commands = []
|
||||
self.commands = {
|
||||
"help": Command(
|
||||
aliases=frozenset(["/help", "/h"]),
|
||||
description="Show help message",
|
||||
handler="_show_help",
|
||||
),
|
||||
"status": Command(
|
||||
aliases=frozenset(["/status", "/stats"]),
|
||||
description="Display agent statistics",
|
||||
handler="_show_status",
|
||||
),
|
||||
"config": Command(
|
||||
aliases=frozenset(["/config", "/cfg", "/theme", "/model"]),
|
||||
description="Edit config settings",
|
||||
handler="_show_config",
|
||||
),
|
||||
"reload": Command(
|
||||
aliases=frozenset(["/reload", "/r"]),
|
||||
description="Reload configuration from disk",
|
||||
handler="_reload_config",
|
||||
),
|
||||
"clear": Command(
|
||||
aliases=frozenset(["/clear", "/reset"]),
|
||||
description="Clear conversation history",
|
||||
handler="_clear_history",
|
||||
),
|
||||
"log": Command(
|
||||
aliases=frozenset(["/log", "/logpath"]),
|
||||
description="Show path to current interaction log file",
|
||||
handler="_show_log_path",
|
||||
),
|
||||
"compact": Command(
|
||||
aliases=frozenset(["/compact", "/summarize"]),
|
||||
description="Compact conversation history by summarizing",
|
||||
handler="_compact_history",
|
||||
),
|
||||
"exit": Command(
|
||||
aliases=frozenset(["/exit", "/quit", "/q"]),
|
||||
description="Exit the application",
|
||||
handler="_exit_app",
|
||||
exits=True,
|
||||
),
|
||||
}
|
||||
|
||||
for command in excluded_commands:
|
||||
self.commands.pop(command, None)
|
||||
|
||||
self._alias_map = {}
|
||||
for cmd_name, cmd in self.commands.items():
|
||||
for alias in cmd.aliases:
|
||||
self._alias_map[alias] = cmd_name
|
||||
|
||||
def find_command(self, user_input: str) -> Command | None:
|
||||
cmd_name = self._alias_map.get(user_input.lower().strip())
|
||||
return self.commands.get(cmd_name) if cmd_name else None
|
||||
|
||||
def get_help_text(self) -> str:
|
||||
lines: list[str] = [
|
||||
"### Keyboard Shortcuts",
|
||||
"",
|
||||
"- `Enter` Submit message",
|
||||
"- `Ctrl+J` / `Shift+Enter` Insert newline",
|
||||
"- `Escape` Interrupt agent or close dialogs",
|
||||
"- `Ctrl+C` Quit (or clear input if text present)",
|
||||
"- `Ctrl+O` Toggle tool output view",
|
||||
"- `Ctrl+T` Toggle todo view",
|
||||
"- `Shift+Tab` Toggle auto-approve mode",
|
||||
"",
|
||||
"### Special Features",
|
||||
"",
|
||||
"- `!<command>` Execute bash command directly",
|
||||
"- `@path/to/file/` Autocompletes file paths",
|
||||
"",
|
||||
"### Commands",
|
||||
"",
|
||||
]
|
||||
|
||||
for cmd in self.commands.values():
|
||||
aliases = ", ".join(f"`{alias}`" for alias in sorted(cmd.aliases))
|
||||
lines.append(f"- {aliases}: {cmd.description}")
|
||||
return "\n".join(lines)
|
||||
261
vibe/cli/entrypoint.py
Normal file
261
vibe/cli/entrypoint.py
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
from rich import print as rprint
|
||||
|
||||
from vibe.cli.textual_ui.app import run_textual_ui
|
||||
from vibe.core.config import (
|
||||
CONFIG_FILE,
|
||||
HISTORY_FILE,
|
||||
INSTRUCTIONS_FILE,
|
||||
MissingAPIKeyError,
|
||||
MissingPromptFileError,
|
||||
VibeConfig,
|
||||
load_api_keys_from_env,
|
||||
)
|
||||
from vibe.core.interaction_logger import InteractionLogger
|
||||
from vibe.core.programmatic import run_programmatic
|
||||
from vibe.core.types import OutputFormat, ResumeSessionInfo
|
||||
from vibe.core.utils import ConversationLimitException
|
||||
from vibe.setup.onboarding import run_onboarding
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run the Mistral Vibe interactive CLI")
|
||||
parser.add_argument(
|
||||
"initial_prompt",
|
||||
nargs="?",
|
||||
metavar="PROMPT",
|
||||
help="Initial prompt to start the interactive session with.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--prompt",
|
||||
nargs="?",
|
||||
const="",
|
||||
metavar="TEXT",
|
||||
help="Run in programmatic mode: send prompt, auto-approve all tools, "
|
||||
"output response, and exit.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-approve",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Automatically approve all tool executions.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-turns",
|
||||
type=int,
|
||||
metavar="N",
|
||||
help="Maximum number of assistant turns "
|
||||
"(only applies in programmatic mode with -p).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-price",
|
||||
type=float,
|
||||
metavar="DOLLARS",
|
||||
help="Maximum cost in dollars (only applies in programmatic mode with -p). "
|
||||
"Session will be interrupted if cost exceeds this limit.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enabled-tools",
|
||||
action="append",
|
||||
metavar="TOOL",
|
||||
help="Enable specific tools. In programmatic mode (-p), this disables "
|
||||
"all other tools. "
|
||||
"Can use exact names, glob patterns (e.g., 'bash*'), or "
|
||||
"regex with 're:' prefix. Can be specified multiple times.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
choices=["text", "json", "streaming"],
|
||||
default="text",
|
||||
help="Output format for programmatic mode (-p): 'text' "
|
||||
"for human-readable (default), 'json' for all messages at end, "
|
||||
"'streaming' for newline-delimited JSON per message.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--agent",
|
||||
metavar="NAME",
|
||||
default=None,
|
||||
help="Load agent configuration from ~/.vibe/agents/NAME.toml",
|
||||
)
|
||||
parser.add_argument("--setup", action="store_true", help="Setup API key and exit")
|
||||
|
||||
continuation_group = parser.add_mutually_exclusive_group()
|
||||
continuation_group.add_argument(
|
||||
"-c",
|
||||
"--continue",
|
||||
action="store_true",
|
||||
dest="continue_session",
|
||||
help="Continue from the most recent saved session",
|
||||
)
|
||||
continuation_group.add_argument(
|
||||
"--resume",
|
||||
metavar="SESSION_ID",
|
||||
help="Resume a specific session by its ID (supports partial matching)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_prompt_from_stdin() -> str | None:
|
||||
if sys.stdin.isatty():
|
||||
return None
|
||||
try:
|
||||
if content := sys.stdin.read().strip():
|
||||
sys.stdin = sys.__stdin__ = open("/dev/tty")
|
||||
return content
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def load_config_or_exit(agent: str | None = None) -> VibeConfig:
|
||||
try:
|
||||
return VibeConfig.load(agent)
|
||||
except MissingAPIKeyError:
|
||||
run_onboarding()
|
||||
return VibeConfig.load(agent)
|
||||
except MissingPromptFileError as e:
|
||||
rprint(f"[yellow]Invalid system prompt id: {e}[/]")
|
||||
sys.exit(1)
|
||||
except ValueError as e:
|
||||
rprint(f"[yellow]{e}[/]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def main() -> None: # noqa: PLR0912, PLR0915
|
||||
load_api_keys_from_env()
|
||||
args = parse_arguments()
|
||||
|
||||
if args.setup:
|
||||
run_onboarding()
|
||||
sys.exit(0)
|
||||
try:
|
||||
if not CONFIG_FILE.exists():
|
||||
try:
|
||||
VibeConfig.save_updates(VibeConfig.create_default())
|
||||
except Exception as e:
|
||||
rprint(f"[yellow]Could not create default config file: {e}[/]")
|
||||
|
||||
if not INSTRUCTIONS_FILE.exists():
|
||||
try:
|
||||
INSTRUCTIONS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
INSTRUCTIONS_FILE.touch()
|
||||
except Exception as e:
|
||||
rprint(f"[yellow]Could not create instructions file: {e}[/]")
|
||||
|
||||
if not HISTORY_FILE.exists():
|
||||
try:
|
||||
HISTORY_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
HISTORY_FILE.write_text("Hello Vibe!\n", "utf-8")
|
||||
except Exception as e:
|
||||
rprint(f"[yellow]Could not create history file: {e}[/]")
|
||||
|
||||
config = load_config_or_exit(args.agent)
|
||||
|
||||
if args.enabled_tools:
|
||||
config.enabled_tools = args.enabled_tools
|
||||
|
||||
loaded_messages = None
|
||||
session_info = None
|
||||
|
||||
if args.continue_session or args.resume:
|
||||
if not config.session_logging.enabled:
|
||||
rprint(
|
||||
"[red]Session logging is disabled. "
|
||||
"Enable it in config to use --continue or --resume[/]"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
session_to_load = None
|
||||
if args.continue_session:
|
||||
session_to_load = InteractionLogger.find_latest_session(
|
||||
config.session_logging
|
||||
)
|
||||
if not session_to_load:
|
||||
rprint(
|
||||
f"[red]No previous sessions found in "
|
||||
f"{config.session_logging.save_dir}[/]"
|
||||
)
|
||||
sys.exit(1)
|
||||
else:
|
||||
session_to_load = InteractionLogger.find_session_by_id(
|
||||
args.resume, config.session_logging
|
||||
)
|
||||
if not session_to_load:
|
||||
rprint(
|
||||
f"[red]Session '{args.resume}' not found in "
|
||||
f"{config.session_logging.save_dir}[/]"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
loaded_messages, metadata = InteractionLogger.load_session(
|
||||
session_to_load
|
||||
)
|
||||
session_id = metadata.get("session_id", "unknown")[:8]
|
||||
session_time = metadata.get("start_time", "unknown time")
|
||||
|
||||
session_info = ResumeSessionInfo(
|
||||
type="continue" if args.continue_session else "resume",
|
||||
session_id=session_id,
|
||||
session_time=session_time,
|
||||
)
|
||||
except Exception as e:
|
||||
rprint(f"[red]Failed to load session: {e}[/]")
|
||||
sys.exit(1)
|
||||
|
||||
stdin_prompt = get_prompt_from_stdin()
|
||||
if args.prompt is not None:
|
||||
programmatic_prompt = args.prompt or stdin_prompt
|
||||
if not programmatic_prompt:
|
||||
print(
|
||||
"Error: No prompt provided for programmatic mode", file=sys.stderr
|
||||
)
|
||||
sys.exit(1)
|
||||
output_format = OutputFormat(
|
||||
args.output if hasattr(args, "output") else "text"
|
||||
)
|
||||
|
||||
try:
|
||||
final_response = run_programmatic(
|
||||
config=config,
|
||||
prompt=programmatic_prompt,
|
||||
max_turns=args.max_turns,
|
||||
max_price=args.max_price,
|
||||
output_format=output_format,
|
||||
previous_messages=loaded_messages,
|
||||
)
|
||||
if final_response:
|
||||
print(final_response)
|
||||
sys.exit(0)
|
||||
except ConversationLimitException as e:
|
||||
print(e, file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except RuntimeError as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
else:
|
||||
run_textual_ui(
|
||||
config,
|
||||
auto_approve=args.auto_approve,
|
||||
enable_streaming=True,
|
||||
initial_prompt=args.initial_prompt or stdin_prompt,
|
||||
loaded_messages=loaded_messages,
|
||||
session_info=session_info,
|
||||
)
|
||||
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
rprint("\n[dim]Bye![/]")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
91
vibe/cli/history_manager.py
Normal file
91
vibe/cli/history_manager.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class HistoryManager:
|
||||
def __init__(self, history_file: Path, max_entries: int = 100) -> None:
|
||||
self.history_file = history_file
|
||||
self.max_entries = max_entries
|
||||
self._entries: list[str] = []
|
||||
self._current_index: int = -1
|
||||
self._temp_input: str = ""
|
||||
self._load_history()
|
||||
|
||||
def _load_history(self) -> None:
|
||||
if not self.history_file.exists():
|
||||
return
|
||||
|
||||
try:
|
||||
with self.history_file.open("r", encoding="utf-8") as f:
|
||||
entries = []
|
||||
for raw_line in f:
|
||||
raw_line = raw_line.rstrip("\n\r")
|
||||
if not raw_line:
|
||||
continue
|
||||
try:
|
||||
entry = json.loads(raw_line)
|
||||
except json.JSONDecodeError:
|
||||
entry = raw_line
|
||||
entries.append(entry if isinstance(entry, str) else str(entry))
|
||||
self._entries = entries[-self.max_entries :]
|
||||
except (OSError, UnicodeDecodeError):
|
||||
self._entries = []
|
||||
|
||||
def _save_history(self) -> None:
|
||||
try:
|
||||
self.history_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self.history_file.open("w", encoding="utf-8") as f:
|
||||
for entry in self._entries:
|
||||
f.write(json.dumps(entry) + "\n")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def add(self, text: str) -> None:
|
||||
text = text.strip()
|
||||
if not text or text.startswith("/"):
|
||||
return
|
||||
|
||||
if self._entries and self._entries[-1] == text:
|
||||
return
|
||||
|
||||
self._entries.append(text)
|
||||
|
||||
if len(self._entries) > self.max_entries:
|
||||
self._entries = self._entries[-self.max_entries :]
|
||||
|
||||
self._save_history()
|
||||
self.reset_navigation()
|
||||
|
||||
def get_previous(self, current_input: str, prefix: str = "") -> str | None:
|
||||
if not self._entries:
|
||||
return None
|
||||
|
||||
if self._current_index == -1:
|
||||
self._temp_input = current_input
|
||||
self._current_index = len(self._entries)
|
||||
|
||||
for i in range(self._current_index - 1, -1, -1):
|
||||
if self._entries[i].startswith(prefix):
|
||||
self._current_index = i
|
||||
return self._entries[i]
|
||||
|
||||
return None
|
||||
|
||||
def get_next(self, prefix: str = "") -> str | None:
|
||||
if self._current_index == -1:
|
||||
return None
|
||||
|
||||
for i in range(self._current_index + 1, len(self._entries)):
|
||||
if self._entries[i].startswith(prefix):
|
||||
self._current_index = i
|
||||
return self._entries[i]
|
||||
|
||||
result = self._temp_input
|
||||
self.reset_navigation()
|
||||
return result
|
||||
|
||||
def reset_navigation(self) -> None:
|
||||
self._current_index = -1
|
||||
self._temp_input = ""
|
||||
0
vibe/cli/textual_ui/__init__.py
Normal file
0
vibe/cli/textual_ui/__init__.py
Normal file
1099
vibe/cli/textual_ui/app.py
Normal file
1099
vibe/cli/textual_ui/app.py
Normal file
File diff suppressed because it is too large
Load diff
681
vibe/cli/textual_ui/app.tcss
Normal file
681
vibe/cli/textual_ui/app.tcss
Normal file
|
|
@ -0,0 +1,681 @@
|
|||
Screen {
|
||||
background: $background;
|
||||
}
|
||||
|
||||
#chat {
|
||||
height: 1fr;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
padding: 0 2 0 2;
|
||||
}
|
||||
|
||||
#loading-area {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
padding: 1 2 0 2;
|
||||
layout: horizontal;
|
||||
align: left middle;
|
||||
}
|
||||
|
||||
#loading-area-content {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
align: left middle;
|
||||
}
|
||||
|
||||
#todo-area {
|
||||
height: auto;
|
||||
max-height: 8;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
overflow-y: auto;
|
||||
text-align: left;
|
||||
margin-top: 0;
|
||||
}
|
||||
|
||||
#bottom-app-container {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
}
|
||||
|
||||
#bottom-bar {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
padding: 0 2 1 2;
|
||||
align: left middle;
|
||||
layout: horizontal;
|
||||
}
|
||||
|
||||
#spacer {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
#messages {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
text-align: left;
|
||||
margin-top: 1;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
#input-container {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
padding: 0;
|
||||
margin: 0 2 1 2;
|
||||
}
|
||||
|
||||
#completion-popup {
|
||||
width: 100%;
|
||||
padding: 1 1 1 1;
|
||||
color: $text;
|
||||
}
|
||||
|
||||
#input-box {
|
||||
height: auto;
|
||||
width: 100%;
|
||||
background: $background;
|
||||
border: round $foreground-muted;
|
||||
padding: 0 1;
|
||||
|
||||
&.border-warning {
|
||||
border: round $warning;
|
||||
}
|
||||
}
|
||||
|
||||
#input-body {
|
||||
height: auto;
|
||||
}
|
||||
|
||||
#prompt {
|
||||
width: auto;
|
||||
background: transparent;
|
||||
color: $primary;
|
||||
text-style: bold;
|
||||
padding: 0 1 0 0;
|
||||
}
|
||||
|
||||
#input {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
max-height: 16;
|
||||
background: transparent;
|
||||
color: $text;
|
||||
border: none;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
ToastRack {
|
||||
align: left bottom;
|
||||
padding: 0 2;
|
||||
margin: 0 2 6 2;
|
||||
}
|
||||
|
||||
Markdown MarkdownFence {
|
||||
overflow-x: auto;
|
||||
scrollbar-size-horizontal: 1;
|
||||
max-width: 95%;
|
||||
}
|
||||
|
||||
.user-message {
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
height: auto;
|
||||
|
||||
&:first-child {
|
||||
margin-top: 0;
|
||||
}
|
||||
|
||||
&.pending {
|
||||
.user-message-prompt,
|
||||
.user-message-content {
|
||||
opacity: 0.7;
|
||||
text-style: italic;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
.user-message-container {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.user-message-prompt {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $primary;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.user-message-content {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.assistant-message {
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.assistant-message-container {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
align: left top;
|
||||
}
|
||||
|
||||
.assistant-message-dot {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.assistant-message-content {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
padding: 0;
|
||||
|
||||
Markdown {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
}
|
||||
}
|
||||
|
||||
.interrupt-message {
|
||||
margin-top: 0;
|
||||
margin-bottom: 0;
|
||||
margin-left: 2;
|
||||
height: auto;
|
||||
padding: 1 2;
|
||||
background: $warning 10%;
|
||||
color: $text-warning;
|
||||
}
|
||||
|
||||
.error-message {
|
||||
margin-top: 1;
|
||||
height: auto;
|
||||
padding: 1 2;
|
||||
background: $error 10%;
|
||||
color: $error;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.bash-output-message {
|
||||
margin-top: 1;
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.bash-output-container {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
padding: 1 2;
|
||||
background: $surface;
|
||||
}
|
||||
|
||||
.bash-cwd-line,
|
||||
.bash-command-line {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
margin-bottom: 1;
|
||||
align: left middle;
|
||||
}
|
||||
|
||||
.bash-cwd {
|
||||
width: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.bash-cwd-spacer,
|
||||
.bash-command-spacer {
|
||||
width: 1fr;
|
||||
}
|
||||
|
||||
.bash-chevron {
|
||||
width: auto;
|
||||
color: $primary;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.bash-command {
|
||||
width: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.bash-output {
|
||||
width: 100%;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.bash-exit-success {
|
||||
width: auto;
|
||||
color: $text-success;
|
||||
}
|
||||
|
||||
.bash-exit-failure {
|
||||
width: auto;
|
||||
color: $text-error;
|
||||
}
|
||||
|
||||
.bash-exit-code {
|
||||
width: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.unknown-event {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
BlinkingMessage {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
|
||||
Horizontal {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
}
|
||||
|
||||
.blink-dot {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
|
||||
&.success {
|
||||
color: $text-success;
|
||||
}
|
||||
|
||||
&.error {
|
||||
color: $text-error;
|
||||
}
|
||||
}
|
||||
|
||||
.blink-text {
|
||||
width: 1fr;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.compact-message {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
margin-top: 1;
|
||||
}
|
||||
|
||||
.tool-call {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
margin-top: 1;
|
||||
}
|
||||
|
||||
.tool-result {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
margin-top: 0;
|
||||
margin-left: 2;
|
||||
padding: 1 2;
|
||||
background: $surface;
|
||||
color: $foreground;
|
||||
|
||||
&.error-text {
|
||||
background: $error 10%;
|
||||
color: $text-error;
|
||||
}
|
||||
|
||||
&.warning-text {
|
||||
background: $warning 10%;
|
||||
color: $text-warning;
|
||||
}
|
||||
}
|
||||
|
||||
.tool-call-widget {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.tool-call-detail {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.tool-result-widget {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
|
||||
Static {
|
||||
height: auto;
|
||||
}
|
||||
}
|
||||
|
||||
.tool-result-detail {
|
||||
height: auto;
|
||||
margin-top: 0;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.tool-result-error {
|
||||
color: $text-error;
|
||||
}
|
||||
|
||||
.tool-result-warning {
|
||||
color: $text-warning;
|
||||
}
|
||||
|
||||
.diff-header {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.diff-removed {
|
||||
height: auto;
|
||||
color: $text-error;
|
||||
}
|
||||
|
||||
.diff-added {
|
||||
height: auto;
|
||||
color: $text-success;
|
||||
}
|
||||
|
||||
.diff-range {
|
||||
height: auto;
|
||||
color: $primary;
|
||||
}
|
||||
|
||||
.diff-context {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.todo-empty {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.todo-pending {
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.todo-in_progress {
|
||||
height: auto;
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
.todo-completed {
|
||||
height: auto;
|
||||
color: $success;
|
||||
}
|
||||
|
||||
.todo-cancelled {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
#todo-area .tool-result {
|
||||
margin-left: 0;
|
||||
background: $surface;
|
||||
}
|
||||
|
||||
.loading-widget {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.loading-container {
|
||||
width: auto;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.loading-star {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
.loading-status {
|
||||
width: auto;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.loading-char {
|
||||
width: auto;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.loading-ellipsis {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.loading-hint {
|
||||
width: auto;
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
WelcomeBanner {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
border: round $surface;
|
||||
border-title-align: center;
|
||||
text-align: center;
|
||||
content-align: center middle;
|
||||
padding: 2 4;
|
||||
margin: 1 1 0 1;
|
||||
color: $foreground;
|
||||
|
||||
.muted {
|
||||
color: $text-muted;
|
||||
}
|
||||
}
|
||||
|
||||
#config-app {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
background: $background;
|
||||
border: round $foreground-muted;
|
||||
padding: 0 1;
|
||||
margin: 0 2 1 2;
|
||||
}
|
||||
|
||||
#config-content {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.settings-title {
|
||||
height: auto;
|
||||
text-style: bold;
|
||||
color: $primary;
|
||||
}
|
||||
|
||||
.settings-option {
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.settings-cursor-selected {
|
||||
color: $primary;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.settings-label-selected {
|
||||
color: $text;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.settings-value-toggle-on-selected {
|
||||
color: $text-success;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.settings-value-toggle-on-unselected {
|
||||
color: $success;
|
||||
}
|
||||
|
||||
.settings-value-toggle-off {
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.settings-value-cycle-selected {
|
||||
color: $primary;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
.settings-value-cycle-unselected {
|
||||
color: $primary;
|
||||
}
|
||||
|
||||
.settings-help {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
#approval-app {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
max-height: 16;
|
||||
background: $background;
|
||||
border: round $foreground-muted;
|
||||
padding: 0 1;
|
||||
margin: 0 2 1 2;
|
||||
}
|
||||
|
||||
#approval-content {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.approval-tool-info-scroll {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
max-height: 8;
|
||||
}
|
||||
|
||||
.approval-title {
|
||||
height: auto;
|
||||
text-style: bold;
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
.approval-tool-info-container {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.tool-approval-widget {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
|
||||
Static {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
Vertical {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
}
|
||||
|
||||
.approval-option {
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.approval-cursor-selected {
|
||||
&.approval-option-yes {
|
||||
color: $text-success;
|
||||
text-style: bold;
|
||||
}
|
||||
|
||||
&.approval-option-no {
|
||||
color: $text-error;
|
||||
text-style: bold;
|
||||
}
|
||||
}
|
||||
|
||||
.approval-option-selected {
|
||||
&.approval-option-yes {
|
||||
color: $success;
|
||||
}
|
||||
|
||||
&.approval-option-no {
|
||||
color: $error;
|
||||
}
|
||||
}
|
||||
|
||||
.approval-help {
|
||||
height: auto;
|
||||
color: $text-muted;
|
||||
}
|
||||
|
||||
.approval-description {
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
}
|
||||
|
||||
.code-block {
|
||||
height: auto;
|
||||
color: $foreground;
|
||||
background: $surface;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
Horizontal {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
ModeIndicator {
|
||||
width: auto;
|
||||
height: auto;
|
||||
background: transparent;
|
||||
padding: 0;
|
||||
margin: 0 0 0 1;
|
||||
color: $warning;
|
||||
align: left middle;
|
||||
|
||||
&.mode-on {
|
||||
color: $warning;
|
||||
}
|
||||
|
||||
&.mode-off {
|
||||
color: $text-muted;
|
||||
}
|
||||
}
|
||||
|
||||
PathDisplay {
|
||||
width: auto;
|
||||
height: auto;
|
||||
background: transparent;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
color: $primary;
|
||||
}
|
||||
|
||||
ContextProgress {
|
||||
width: auto;
|
||||
height: auto;
|
||||
background: transparent;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
color: $foreground;
|
||||
}
|
||||
5
vibe/cli/textual_ui/handlers/__init__.py
Normal file
5
vibe/cli/textual_ui/handlers/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.textual_ui.handlers.event_handler import EventHandler
|
||||
|
||||
__all__ = ["EventHandler"]
|
||||
158
vibe/cli/textual_ui/handlers/event_handler.py
Normal file
158
vibe/cli/textual_ui/handlers/event_handler.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from textual.widgets import Static
|
||||
|
||||
from vibe.cli.textual_ui.widgets.compact import CompactMessage
|
||||
from vibe.cli.textual_ui.widgets.messages import AssistantMessage
|
||||
from vibe.cli.textual_ui.widgets.tools import ToolCallMessage, ToolResultMessage
|
||||
from vibe.core.types import (
|
||||
AssistantEvent,
|
||||
BaseEvent,
|
||||
CompactEndEvent,
|
||||
CompactStartEvent,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
)
|
||||
from vibe.core.utils import TaggedText
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.cli.textual_ui.widgets.loading import LoadingWidget
|
||||
|
||||
|
||||
class EventHandler:
|
||||
def __init__(
|
||||
self,
|
||||
mount_callback: Callable,
|
||||
scroll_callback: Callable,
|
||||
todo_area_callback: Callable,
|
||||
get_tools_collapsed: Callable[[], bool],
|
||||
get_todos_collapsed: Callable[[], bool],
|
||||
) -> None:
|
||||
self.mount_callback = mount_callback
|
||||
self.scroll_callback = scroll_callback
|
||||
self.todo_area_callback = todo_area_callback
|
||||
self.get_tools_collapsed = get_tools_collapsed
|
||||
self.get_todos_collapsed = get_todos_collapsed
|
||||
self.current_tool_call: ToolCallMessage | None = None
|
||||
self.current_compact: CompactMessage | None = None
|
||||
self.tool_results: list[ToolResultMessage] = []
|
||||
|
||||
async def handle_event(
|
||||
self,
|
||||
event: BaseEvent,
|
||||
loading_active: bool = False,
|
||||
loading_widget: LoadingWidget | None = None,
|
||||
) -> ToolCallMessage | None:
|
||||
match event:
|
||||
case ToolCallEvent():
|
||||
return await self._handle_tool_call(event, loading_widget)
|
||||
case ToolResultEvent():
|
||||
sanitized_event = self._sanitize_event(event)
|
||||
|
||||
await self._handle_tool_result(sanitized_event)
|
||||
return None
|
||||
case AssistantEvent():
|
||||
await self._handle_assistant_message(event)
|
||||
return None
|
||||
case CompactStartEvent():
|
||||
await self._handle_compact_start()
|
||||
return None
|
||||
case CompactEndEvent():
|
||||
await self._handle_compact_end(event)
|
||||
return None
|
||||
case _:
|
||||
await self._handle_unknown_event(event)
|
||||
return None
|
||||
|
||||
def _sanitize_event(self, event: ToolResultEvent) -> ToolResultEvent:
|
||||
if isinstance(event, ToolResultEvent):
|
||||
return ToolResultEvent(
|
||||
tool_name=event.tool_name,
|
||||
tool_class=event.tool_class,
|
||||
result=event.result,
|
||||
error=TaggedText.from_string(event.error).message
|
||||
if event.error
|
||||
else None,
|
||||
skipped=event.skipped,
|
||||
skip_reason=TaggedText.from_string(event.skip_reason).message
|
||||
if event.skip_reason
|
||||
else None,
|
||||
duration=event.duration,
|
||||
tool_call_id=event.tool_call_id,
|
||||
)
|
||||
return event
|
||||
|
||||
async def _handle_tool_call(
|
||||
self, event: ToolCallEvent, loading_widget: LoadingWidget | None = None
|
||||
) -> ToolCallMessage | None:
|
||||
tool_call = ToolCallMessage(event)
|
||||
|
||||
if loading_widget and event.tool_class:
|
||||
from vibe.core.tools.ui import ToolUIDataAdapter
|
||||
|
||||
adapter = ToolUIDataAdapter(event.tool_class)
|
||||
status_text = adapter.get_status_text()
|
||||
loading_widget.set_status(status_text)
|
||||
|
||||
# Don't show todo in messages
|
||||
if event.tool_name != "todo":
|
||||
await self.mount_callback(tool_call)
|
||||
|
||||
self.current_tool_call = tool_call
|
||||
return tool_call
|
||||
|
||||
async def _handle_tool_result(self, event: ToolResultEvent) -> None:
|
||||
if event.tool_name == "todo":
|
||||
todos_collapsed = self.get_todos_collapsed()
|
||||
tool_result = ToolResultMessage(
|
||||
event, self.current_tool_call, collapsed=todos_collapsed
|
||||
)
|
||||
# Show in todo area
|
||||
todo_area = self.todo_area_callback()
|
||||
await todo_area.remove_children()
|
||||
await todo_area.mount(tool_result)
|
||||
else:
|
||||
tools_collapsed = self.get_tools_collapsed()
|
||||
tool_result = ToolResultMessage(
|
||||
event, self.current_tool_call, collapsed=tools_collapsed
|
||||
)
|
||||
await self.mount_callback(tool_result)
|
||||
|
||||
self.tool_results.append(tool_result)
|
||||
self.current_tool_call = None
|
||||
|
||||
async def _handle_assistant_message(self, event: AssistantEvent) -> None:
|
||||
await self.mount_callback(AssistantMessage(event.content))
|
||||
|
||||
async def _handle_compact_start(self) -> None:
|
||||
compact_msg = CompactMessage()
|
||||
self.current_compact = compact_msg
|
||||
await self.mount_callback(compact_msg)
|
||||
|
||||
async def _handle_compact_end(self, event: CompactEndEvent) -> None:
|
||||
if self.current_compact:
|
||||
self.current_compact.set_complete(
|
||||
old_tokens=event.old_context_tokens, new_tokens=event.new_context_tokens
|
||||
)
|
||||
self.current_compact = None
|
||||
|
||||
async def _handle_unknown_event(self, event: BaseEvent) -> None:
|
||||
await self.mount_callback(
|
||||
Static(str(event), markup=False, classes="unknown-event")
|
||||
)
|
||||
|
||||
def stop_current_tool_call(self) -> None:
|
||||
if self.current_tool_call:
|
||||
self.current_tool_call.stop_blinking()
|
||||
self.current_tool_call = None
|
||||
|
||||
def stop_current_compact(self) -> None:
|
||||
if self.current_compact:
|
||||
self.current_compact.stop_blinking(success=False)
|
||||
self.current_compact = None
|
||||
|
||||
def get_last_tool_result(self) -> ToolResultMessage | None:
|
||||
return self.tool_results[-1] if self.tool_results else None
|
||||
5
vibe/cli/textual_ui/renderers/__init__.py
Normal file
5
vibe/cli/textual_ui/renderers/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.textual_ui.renderers.tool_renderers import get_renderer
|
||||
|
||||
__all__ = ["get_renderer"]
|
||||
216
vibe/cli/textual_ui/renderers/tool_renderers.py
Normal file
216
vibe/cli/textual_ui/renderers/tool_renderers.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import difflib
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.tools.ui import ToolResultDisplay
|
||||
|
||||
from vibe.cli.textual_ui.widgets.tool_widgets import (
|
||||
BashApprovalWidget,
|
||||
BashResultWidget,
|
||||
GrepApprovalWidget,
|
||||
GrepResultWidget,
|
||||
ReadFileApprovalWidget,
|
||||
ReadFileResultWidget,
|
||||
SearchReplaceApprovalWidget,
|
||||
SearchReplaceResultWidget,
|
||||
TodoApprovalWidget,
|
||||
TodoResultWidget,
|
||||
ToolApprovalWidget,
|
||||
ToolResultWidget,
|
||||
WriteFileApprovalWidget,
|
||||
WriteFileResultWidget,
|
||||
)
|
||||
|
||||
|
||||
class ToolRenderer:
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[ToolApprovalWidget], dict[str, Any]]:
|
||||
return ToolApprovalWidget, tool_args
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[ToolResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"details": self._clean_details(display.details),
|
||||
"warnings": display.warnings,
|
||||
}
|
||||
return ToolResultWidget, data
|
||||
|
||||
def _clean_details(self, details: dict) -> dict:
|
||||
clean = {}
|
||||
for key, value in details.items():
|
||||
if value is None or value in ("", []):
|
||||
continue
|
||||
value_str = str(value).strip().replace("\n", " ").replace("\r", "")
|
||||
value_str = " ".join(value_str.split())
|
||||
if value_str:
|
||||
clean[key] = value_str
|
||||
return clean
|
||||
|
||||
|
||||
class BashRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[BashApprovalWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"command": tool_args.get("command", ""),
|
||||
"description": tool_args.get("description", ""),
|
||||
}
|
||||
return BashApprovalWidget, data
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[BashResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"details": self._clean_details(display.details),
|
||||
"warnings": display.warnings,
|
||||
}
|
||||
return BashResultWidget, data
|
||||
|
||||
|
||||
class WriteFileRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[WriteFileApprovalWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"path": tool_args.get("path", ""),
|
||||
"content": tool_args.get("content", ""),
|
||||
"file_extension": tool_args.get("file_extension", "text"),
|
||||
}
|
||||
return WriteFileApprovalWidget, data
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[WriteFileResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"path": display.details.get("path", ""),
|
||||
"bytes_written": display.details.get("bytes_written"),
|
||||
"content": display.details.get("content", ""),
|
||||
"file_extension": display.details.get("file_extension", "text"),
|
||||
}
|
||||
return WriteFileResultWidget, data
|
||||
|
||||
|
||||
class SearchReplaceRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[SearchReplaceApprovalWidget], dict[str, Any]]:
|
||||
file_path = tool_args.get("file_path", "")
|
||||
content = str(tool_args.get("content", ""))
|
||||
|
||||
diff_lines = self._parse_search_replace_blocks(content)
|
||||
|
||||
data = {"file_path": file_path, "diff_lines": diff_lines}
|
||||
return SearchReplaceApprovalWidget, data
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[SearchReplaceResultWidget], dict[str, Any]]:
|
||||
diff_lines = self._parse_search_replace_blocks(
|
||||
display.details.get("content", "")
|
||||
)
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"diff_lines": diff_lines if not collapsed else [],
|
||||
}
|
||||
return SearchReplaceResultWidget, data
|
||||
|
||||
def _parse_search_replace_blocks(self, content: str) -> list[str]:
|
||||
if "<<<<<<< SEARCH" not in content:
|
||||
return [content]
|
||||
|
||||
try:
|
||||
sections = content.split("<<<<<<< SEARCH")
|
||||
rest = sections[1].split("=======")
|
||||
search_section = rest[0].strip()
|
||||
replace_part = rest[1].split(">>>>>>> REPLACE")
|
||||
replace_section = replace_part[0].strip()
|
||||
|
||||
search_lines = search_section.split("\n")
|
||||
replace_lines = replace_section.split("\n")
|
||||
|
||||
diff = difflib.unified_diff(search_lines, replace_lines, lineterm="", n=2)
|
||||
return list(diff)[2:] # Skip file headers
|
||||
except (IndexError, AttributeError):
|
||||
return [content[:500]]
|
||||
|
||||
|
||||
class TodoRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[TodoApprovalWidget], dict[str, Any]]:
|
||||
data = {"description": tool_args.get("description", "")}
|
||||
return TodoApprovalWidget, data
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[TodoResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"todos_by_status": display.details.get("todos_by_status", {}),
|
||||
}
|
||||
return TodoResultWidget, data
|
||||
|
||||
|
||||
class ReadFileRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[ReadFileApprovalWidget], dict[str, Any]]:
|
||||
return ReadFileApprovalWidget, tool_args
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[ReadFileResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"path": display.details.get("path", ""),
|
||||
"warnings": display.warnings,
|
||||
"content": display.details.get("content", "") if not collapsed else "",
|
||||
"file_extension": display.details.get("file_extension", "text"),
|
||||
}
|
||||
return ReadFileResultWidget, data
|
||||
|
||||
|
||||
class GrepRenderer(ToolRenderer):
|
||||
def get_approval_widget(
|
||||
self, tool_args: dict
|
||||
) -> tuple[type[GrepApprovalWidget], dict[str, Any]]:
|
||||
return GrepApprovalWidget, tool_args
|
||||
|
||||
def get_result_widget(
|
||||
self, display: ToolResultDisplay, collapsed: bool
|
||||
) -> tuple[type[GrepResultWidget], dict[str, Any]]:
|
||||
data = {
|
||||
"success": display.success,
|
||||
"message": display.message,
|
||||
"warnings": display.warnings,
|
||||
"matches": display.details.get("matches", "") if not collapsed else "",
|
||||
}
|
||||
return GrepResultWidget, data
|
||||
|
||||
|
||||
_RENDERER_REGISTRY: dict[str, type[ToolRenderer]] = {
|
||||
"write_file": WriteFileRenderer,
|
||||
"search_replace": SearchReplaceRenderer,
|
||||
"todo": TodoRenderer,
|
||||
"read_file": ReadFileRenderer,
|
||||
"bash": BashRenderer,
|
||||
"grep": GrepRenderer,
|
||||
}
|
||||
|
||||
|
||||
def get_renderer(tool_name: str) -> ToolRenderer:
|
||||
renderer_class = _RENDERER_REGISTRY.get(tool_name, ToolRenderer)
|
||||
return renderer_class()
|
||||
0
vibe/cli/textual_ui/widgets/__init__.py
Normal file
0
vibe/cli/textual_ui/widgets/__init__.py
Normal file
196
vibe/cli/textual_ui/widgets/approval_app.py
Normal file
196
vibe/cli/textual_ui/widgets/approval_app.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import ClassVar
|
||||
|
||||
from textual import events
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding, BindingType
|
||||
from textual.containers import Container, Vertical, VerticalScroll
|
||||
from textual.message import Message
|
||||
from textual.widgets import Static
|
||||
|
||||
from vibe.cli.textual_ui.renderers import get_renderer
|
||||
from vibe.core.config import VibeConfig
|
||||
|
||||
|
||||
class ApprovalApp(Container):
|
||||
can_focus = True
|
||||
can_focus_children = False
|
||||
|
||||
BINDINGS: ClassVar[list[BindingType]] = [
|
||||
Binding("up", "move_up", "Up", show=False),
|
||||
Binding("down", "move_down", "Down", show=False),
|
||||
Binding("enter", "select", "Select", show=False),
|
||||
Binding("1", "select_1", "Yes", show=False),
|
||||
Binding("y", "select_1", "Yes", show=False),
|
||||
Binding("2", "select_2", "Always Tool Session", show=False),
|
||||
Binding("3", "select_3", "No", show=False),
|
||||
Binding("n", "select_3", "No", show=False),
|
||||
]
|
||||
|
||||
class ApprovalGranted(Message):
|
||||
def __init__(self, tool_name: str, tool_args: dict) -> None:
|
||||
super().__init__()
|
||||
self.tool_name = tool_name
|
||||
self.tool_args = tool_args
|
||||
|
||||
class ApprovalGrantedAlwaysTool(Message):
|
||||
def __init__(
|
||||
self, tool_name: str, tool_args: dict, save_permanently: bool
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.tool_name = tool_name
|
||||
self.tool_args = tool_args
|
||||
self.save_permanently = save_permanently
|
||||
|
||||
class ApprovalRejected(Message):
|
||||
def __init__(self, tool_name: str, tool_args: dict) -> None:
|
||||
super().__init__()
|
||||
self.tool_name = tool_name
|
||||
self.tool_args = tool_args
|
||||
|
||||
def __init__(
|
||||
self, tool_name: str, tool_args: dict, workdir: str, config: VibeConfig
|
||||
) -> None:
|
||||
super().__init__(id="approval-app")
|
||||
self.tool_name = tool_name
|
||||
self.tool_args = tool_args
|
||||
self.workdir = workdir
|
||||
self.config = config
|
||||
self.selected_option = 0
|
||||
self.content_container: Vertical | None = None
|
||||
self.title_widget: Static | None = None
|
||||
self.tool_info_container: Vertical | None = None
|
||||
self.option_widgets: list[Static] = []
|
||||
self.help_widget: Static | None = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Vertical(id="approval-content"):
|
||||
self.title_widget = Static(
|
||||
f"⚠ {self.tool_name} command", classes="approval-title"
|
||||
)
|
||||
yield self.title_widget
|
||||
|
||||
with VerticalScroll(classes="approval-tool-info-scroll"):
|
||||
self.tool_info_container = Vertical(
|
||||
classes="approval-tool-info-container"
|
||||
)
|
||||
yield self.tool_info_container
|
||||
|
||||
yield Static("")
|
||||
|
||||
for _ in range(3):
|
||||
widget = Static("", classes="approval-option")
|
||||
self.option_widgets.append(widget)
|
||||
yield widget
|
||||
|
||||
yield Static("")
|
||||
|
||||
self.help_widget = Static(
|
||||
"↑↓ navigate Enter select ESC reject", classes="approval-help"
|
||||
)
|
||||
yield self.help_widget
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
await self._update_tool_info()
|
||||
self._update_options()
|
||||
self.focus()
|
||||
|
||||
async def _update_tool_info(self) -> None:
|
||||
if not self.tool_info_container:
|
||||
return
|
||||
|
||||
renderer = get_renderer(self.tool_name)
|
||||
widget_class, data = renderer.get_approval_widget(self.tool_args)
|
||||
|
||||
await self.tool_info_container.remove_children()
|
||||
approval_widget = widget_class(data)
|
||||
await self.tool_info_container.mount(approval_widget)
|
||||
|
||||
def _update_options(self) -> None:
|
||||
options = [
|
||||
("Yes", "yes"),
|
||||
(f"Yes and always allow {self.tool_name} this session", "yes"),
|
||||
("No and tell the agent what to do instead", "no"),
|
||||
]
|
||||
|
||||
for idx, ((text, color_type), widget) in enumerate(
|
||||
zip(options, self.option_widgets, strict=True)
|
||||
):
|
||||
is_selected = idx == self.selected_option
|
||||
|
||||
cursor = "› " if is_selected else " "
|
||||
option_text = f"{cursor}{idx + 1}. {text}"
|
||||
|
||||
widget.update(option_text)
|
||||
|
||||
widget.remove_class("approval-cursor-selected")
|
||||
widget.remove_class("approval-option-selected")
|
||||
widget.remove_class("approval-option-yes")
|
||||
widget.remove_class("approval-option-no")
|
||||
|
||||
if is_selected:
|
||||
widget.add_class("approval-cursor-selected")
|
||||
if color_type == "yes":
|
||||
widget.add_class("approval-option-yes")
|
||||
else:
|
||||
widget.add_class("approval-option-no")
|
||||
else:
|
||||
widget.add_class("approval-option-selected")
|
||||
if color_type == "yes":
|
||||
widget.add_class("approval-option-yes")
|
||||
else:
|
||||
widget.add_class("approval-option-no")
|
||||
|
||||
def action_move_up(self) -> None:
|
||||
self.selected_option = (self.selected_option - 1) % 3
|
||||
self._update_options()
|
||||
|
||||
def action_move_down(self) -> None:
|
||||
self.selected_option = (self.selected_option + 1) % 3
|
||||
self._update_options()
|
||||
|
||||
def action_select(self) -> None:
|
||||
self._handle_selection(self.selected_option)
|
||||
|
||||
def action_select_1(self) -> None:
|
||||
self.selected_option = 0
|
||||
self._handle_selection(0)
|
||||
|
||||
def action_select_2(self) -> None:
|
||||
self.selected_option = 1
|
||||
self._handle_selection(1)
|
||||
|
||||
def action_select_3(self) -> None:
|
||||
self.selected_option = 2
|
||||
self._handle_selection(2)
|
||||
|
||||
def action_reject(self) -> None:
|
||||
self.selected_option = 2
|
||||
self._handle_selection(2)
|
||||
|
||||
def _handle_selection(self, option: int) -> None:
|
||||
match option:
|
||||
case 0:
|
||||
self.post_message(
|
||||
self.ApprovalGranted(
|
||||
tool_name=self.tool_name, tool_args=self.tool_args
|
||||
)
|
||||
)
|
||||
case 1:
|
||||
self.post_message(
|
||||
self.ApprovalGrantedAlwaysTool(
|
||||
tool_name=self.tool_name,
|
||||
tool_args=self.tool_args,
|
||||
save_permanently=False,
|
||||
)
|
||||
)
|
||||
case 2:
|
||||
self.post_message(
|
||||
self.ApprovalRejected(
|
||||
tool_name=self.tool_name, tool_args=self.tool_args
|
||||
)
|
||||
)
|
||||
|
||||
def on_blur(self, event: events.Blur) -> None:
|
||||
self.call_after_refresh(self.focus)
|
||||
67
vibe/cli/textual_ui/widgets/blinking_message.py
Normal file
67
vibe/cli/textual_ui/widgets/blinking_message.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
class BlinkingMessage(Static):
|
||||
def __init__(self, initial_text: str = "", **kwargs: Any) -> None:
|
||||
self.blink_state = False
|
||||
self._blink_timer = None
|
||||
self._is_blinking = True
|
||||
self.success = True
|
||||
self._initial_text = initial_text
|
||||
self._dot_widget: Static | None = None
|
||||
self._text_widget: Static | None = None
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal():
|
||||
self._dot_widget = Static("● ", classes="blink-dot")
|
||||
yield self._dot_widget
|
||||
self._text_widget = Static("", markup=False, classes="blink-text")
|
||||
yield self._text_widget
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self.update_display()
|
||||
self._blink_timer = self.set_interval(0.5, self.toggle_blink)
|
||||
|
||||
def toggle_blink(self) -> None:
|
||||
if not self._is_blinking:
|
||||
return
|
||||
self.blink_state = not self.blink_state
|
||||
self.update_display()
|
||||
|
||||
def update_display(self) -> None:
|
||||
if not self._dot_widget or not self._text_widget:
|
||||
return
|
||||
|
||||
content = self.get_content()
|
||||
|
||||
if self._is_blinking:
|
||||
dot = "● " if self.blink_state else "○ "
|
||||
self._dot_widget.update(dot)
|
||||
self._dot_widget.remove_class("success")
|
||||
self._dot_widget.remove_class("error")
|
||||
else:
|
||||
self._dot_widget.update("● ")
|
||||
if self.success:
|
||||
self._dot_widget.add_class("success")
|
||||
self._dot_widget.remove_class("error")
|
||||
else:
|
||||
self._dot_widget.add_class("error")
|
||||
self._dot_widget.remove_class("success")
|
||||
|
||||
self._text_widget.update(content)
|
||||
|
||||
def get_content(self) -> str:
|
||||
return self._initial_text
|
||||
|
||||
def stop_blinking(self, success: bool = True) -> None:
|
||||
self._is_blinking = False
|
||||
self.blink_state = True
|
||||
self.success = success
|
||||
self.update_display()
|
||||
7
vibe/cli/textual_ui/widgets/chat_input/__init__.py
Normal file
7
vibe/cli/textual_ui/widgets/chat_input/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.textual_ui.widgets.chat_input.body import ChatInputBody
|
||||
from vibe.cli.textual_ui.widgets.chat_input.container import ChatInputContainer
|
||||
from vibe.cli.textual_ui.widgets.chat_input.text_area import ChatTextArea
|
||||
|
||||
__all__ = ["ChatInputBody", "ChatInputContainer", "ChatTextArea"]
|
||||
194
vibe/cli/textual_ui/widgets/chat_input/body.py
Normal file
194
vibe/cli/textual_ui/widgets/chat_input/body.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal
|
||||
from textual.message import Message
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static
|
||||
|
||||
from vibe.cli.history_manager import HistoryManager
|
||||
from vibe.cli.textual_ui.widgets.chat_input.text_area import ChatTextArea
|
||||
|
||||
|
||||
class ChatInputBody(Widget):
|
||||
class Submitted(Message):
|
||||
def __init__(self, value: str) -> None:
|
||||
self.value = value
|
||||
super().__init__()
|
||||
|
||||
def __init__(self, history_file: Path | None = None, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.input_widget: ChatTextArea | None = None
|
||||
self.prompt_widget: Static | None = None
|
||||
|
||||
if history_file:
|
||||
self.history = HistoryManager(history_file)
|
||||
else:
|
||||
self.history = None
|
||||
|
||||
self._completion_reset: Callable[[], None] | None = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal():
|
||||
self.prompt_widget = Static(">", id="prompt")
|
||||
yield self.prompt_widget
|
||||
|
||||
self.input_widget = ChatTextArea(placeholder="Ask anything...", id="input")
|
||||
yield self.input_widget
|
||||
|
||||
def on_mount(self) -> None:
|
||||
if self.input_widget:
|
||||
self.input_widget.focus()
|
||||
|
||||
def _update_prompt(self) -> None:
|
||||
if not self.input_widget or not self.prompt_widget:
|
||||
return
|
||||
|
||||
text = self.input_widget.text
|
||||
if text.startswith("!"):
|
||||
self.prompt_widget.update("!")
|
||||
elif text.startswith("/"):
|
||||
self.prompt_widget.update("/")
|
||||
else:
|
||||
self.prompt_widget.update(">")
|
||||
|
||||
def _load_history_entry(self, text: str, cursor_col: int | None = None) -> None:
|
||||
if not self.input_widget:
|
||||
return
|
||||
|
||||
self.input_widget._navigating_history = True
|
||||
self.input_widget.load_text(text)
|
||||
|
||||
first_line = text.split("\n")[0] if text else ""
|
||||
col = cursor_col if cursor_col is not None else len(first_line)
|
||||
cursor_pos = (0, col)
|
||||
|
||||
self.input_widget.move_cursor(cursor_pos)
|
||||
self.input_widget._last_cursor_col = col
|
||||
self.input_widget._cursor_pos_after_load = cursor_pos
|
||||
self.input_widget._cursor_moved_since_load = False
|
||||
|
||||
self._update_prompt()
|
||||
self._notify_completion_reset()
|
||||
|
||||
def on_chat_text_area_history_previous(
|
||||
self, event: ChatTextArea.HistoryPrevious
|
||||
) -> None:
|
||||
if not self.history or not self.input_widget:
|
||||
return
|
||||
|
||||
if self.history._current_index == -1:
|
||||
self.input_widget._original_text = self.input_widget.text
|
||||
|
||||
if (
|
||||
self.history._current_index != -1
|
||||
and self.input_widget._last_used_prefix is not None
|
||||
and self.input_widget._last_used_prefix != event.prefix
|
||||
):
|
||||
self.history.reset_navigation()
|
||||
|
||||
self.input_widget._last_used_prefix = event.prefix
|
||||
previous = self.history.get_previous(
|
||||
self.input_widget._original_text, prefix=event.prefix
|
||||
)
|
||||
|
||||
if previous is not None:
|
||||
self._load_history_entry(previous)
|
||||
|
||||
def on_chat_text_area_history_next(self, event: ChatTextArea.HistoryNext) -> None:
|
||||
if not self.history or not self.input_widget:
|
||||
return
|
||||
|
||||
if self.history._current_index == -1:
|
||||
return
|
||||
|
||||
if (
|
||||
self.input_widget._last_used_prefix is not None
|
||||
and self.input_widget._last_used_prefix != event.prefix
|
||||
):
|
||||
self.history.reset_navigation()
|
||||
|
||||
self.input_widget._last_used_prefix = event.prefix
|
||||
|
||||
has_next = any(
|
||||
self.history._entries[i].startswith(event.prefix)
|
||||
for i in range(self.history._current_index + 1, len(self.history._entries))
|
||||
)
|
||||
|
||||
original_matches = self.input_widget._original_text.startswith(event.prefix)
|
||||
|
||||
if has_next or original_matches:
|
||||
next_entry = self.history.get_next(prefix=event.prefix)
|
||||
if next_entry is not None:
|
||||
cursor_col = (
|
||||
len(event.prefix) if self.history._current_index == -1 else None
|
||||
)
|
||||
self._load_history_entry(next_entry, cursor_col=cursor_col)
|
||||
|
||||
def on_chat_text_area_history_reset(self, event: ChatTextArea.HistoryReset) -> None:
|
||||
if self.history:
|
||||
self.history.reset_navigation()
|
||||
if self.input_widget:
|
||||
self.input_widget._original_text = ""
|
||||
self.input_widget._cursor_pos_after_load = None
|
||||
self.input_widget._cursor_moved_since_load = False
|
||||
|
||||
def on_text_area_changed(self, event: ChatTextArea.Changed) -> None:
|
||||
self._update_prompt()
|
||||
|
||||
def on_chat_text_area_submitted(self, event: ChatTextArea.Submitted) -> None:
|
||||
event.stop()
|
||||
|
||||
if not self.input_widget:
|
||||
return
|
||||
|
||||
value = event.value.strip()
|
||||
if value:
|
||||
if self.history:
|
||||
self.history.add(value)
|
||||
self.history.reset_navigation()
|
||||
|
||||
self.input_widget.clear_text()
|
||||
self._update_prompt()
|
||||
|
||||
self._notify_completion_reset()
|
||||
|
||||
self.post_message(self.Submitted(value))
|
||||
|
||||
@property
|
||||
def value(self) -> str:
|
||||
return self.input_widget.text if self.input_widget else ""
|
||||
|
||||
@value.setter
|
||||
def value(self, text: str) -> None:
|
||||
if self.input_widget:
|
||||
self.input_widget.load_text(text)
|
||||
self._update_prompt()
|
||||
|
||||
def focus_input(self) -> None:
|
||||
if self.input_widget:
|
||||
self.input_widget.focus()
|
||||
|
||||
def set_completion_reset_callback(
|
||||
self, callback: Callable[[], None] | None
|
||||
) -> None:
|
||||
self._completion_reset = callback
|
||||
|
||||
def _notify_completion_reset(self) -> None:
|
||||
if self._completion_reset:
|
||||
self._completion_reset()
|
||||
|
||||
def replace_input(self, text: str, cursor_offset: int | None = None) -> None:
|
||||
if not self.input_widget:
|
||||
return
|
||||
|
||||
self.input_widget.load_text(text)
|
||||
self.input_widget.reset_history_state()
|
||||
self._update_prompt()
|
||||
|
||||
if cursor_offset is not None:
|
||||
self.input_widget.set_cursor_offset(max(0, min(cursor_offset, len(text))))
|
||||
58
vibe/cli/textual_ui/widgets/chat_input/completion_manager.py
Normal file
58
vibe/cli/textual_ui/widgets/chat_input/completion_manager.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Protocol
|
||||
|
||||
from textual import events
|
||||
|
||||
from vibe.cli.autocompletion.base import CompletionResult
|
||||
|
||||
|
||||
class CompletionController(Protocol):
|
||||
def can_handle(self, text: str, cursor_index: int) -> bool: ...
|
||||
|
||||
def on_text_changed(self, text: str, cursor_index: int) -> None: ...
|
||||
|
||||
def on_key(
|
||||
self, event: events.Key, text: str, cursor_index: int
|
||||
) -> CompletionResult: ...
|
||||
|
||||
def reset(self) -> None: ...
|
||||
|
||||
|
||||
class MultiCompletionManager:
|
||||
def __init__(self, controllers: Sequence[CompletionController]) -> None:
|
||||
self._controllers = list(controllers)
|
||||
self._active: CompletionController | None = None
|
||||
|
||||
def on_text_changed(self, text: str, cursor_index: int) -> None:
|
||||
candidate = None
|
||||
for controller in self._controllers:
|
||||
if controller.can_handle(text, cursor_index):
|
||||
candidate = controller
|
||||
break
|
||||
|
||||
if candidate is None:
|
||||
if self._active is not None:
|
||||
self._active.reset()
|
||||
self._active = None
|
||||
return
|
||||
|
||||
if candidate is not self._active:
|
||||
if self._active is not None:
|
||||
self._active.reset()
|
||||
self._active = candidate
|
||||
|
||||
candidate.on_text_changed(text, cursor_index)
|
||||
|
||||
def on_key(
|
||||
self, event: events.Key, text: str, cursor_index: int
|
||||
) -> CompletionResult:
|
||||
if self._active is None:
|
||||
return CompletionResult.IGNORED
|
||||
return self._active.on_key(event, text, cursor_index)
|
||||
|
||||
def reset(self) -> None:
|
||||
if self._active is not None:
|
||||
self._active.reset()
|
||||
self._active = None
|
||||
43
vibe/cli/textual_ui/widgets/chat_input/completion_popup.py
Normal file
43
vibe/cli/textual_ui/widgets/chat_input/completion_popup.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from rich.text import Text
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
class CompletionPopup(Static):
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
super().__init__("", id="completion-popup", **kwargs)
|
||||
self.styles.display = "none"
|
||||
self.can_focus = False
|
||||
|
||||
def update_suggestions(
|
||||
self, suggestions: list[tuple[str, str]], selected: int
|
||||
) -> None:
|
||||
if not suggestions:
|
||||
self.hide()
|
||||
return
|
||||
|
||||
text = Text()
|
||||
for idx, (label, description) in enumerate(suggestions):
|
||||
if idx:
|
||||
text.append("\n")
|
||||
|
||||
label_style = "bold reverse" if idx == selected else "bold"
|
||||
description_style = "italic" if idx == selected else "dim"
|
||||
|
||||
text.append(label, style=label_style)
|
||||
if description:
|
||||
text.append(" ")
|
||||
text.append(description, style=description_style)
|
||||
|
||||
self.update(text)
|
||||
self.show()
|
||||
|
||||
def hide(self) -> None:
|
||||
self.update("")
|
||||
self.styles.display = "none"
|
||||
|
||||
def show(self) -> None:
|
||||
self.styles.display = "block"
|
||||
157
vibe/cli/textual_ui/widgets/chat_input/container.py
Normal file
157
vibe/cli/textual_ui/widgets/chat_input/container.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Vertical
|
||||
from textual.message import Message
|
||||
|
||||
from vibe.cli.autocompletion.path_completion import PathCompletionController
|
||||
from vibe.cli.autocompletion.slash_command import SlashCommandController
|
||||
from vibe.cli.commands import CommandRegistry
|
||||
from vibe.cli.textual_ui.widgets.chat_input.body import ChatInputBody
|
||||
from vibe.cli.textual_ui.widgets.chat_input.completion_manager import (
|
||||
MultiCompletionManager,
|
||||
)
|
||||
from vibe.cli.textual_ui.widgets.chat_input.completion_popup import CompletionPopup
|
||||
from vibe.cli.textual_ui.widgets.chat_input.text_area import ChatTextArea
|
||||
from vibe.core.autocompletion.completers import CommandCompleter, PathCompleter
|
||||
|
||||
|
||||
class ChatInputContainer(Vertical):
|
||||
ID_INPUT_BOX = "input-box"
|
||||
BORDER_WARNING_CLASS = "border-warning"
|
||||
|
||||
class Submitted(Message):
|
||||
def __init__(self, value: str) -> None:
|
||||
self.value = value
|
||||
super().__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
history_file: Path | None = None,
|
||||
command_registry: CommandRegistry | None = None,
|
||||
show_warning: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self._history_file = history_file
|
||||
self._command_registry = command_registry or CommandRegistry()
|
||||
self._show_warning = show_warning
|
||||
|
||||
command_entries = [
|
||||
(alias, command.description)
|
||||
for command in self._command_registry.commands.values()
|
||||
for alias in sorted(command.aliases)
|
||||
]
|
||||
|
||||
self._completion_manager = MultiCompletionManager([
|
||||
SlashCommandController(CommandCompleter(command_entries), self),
|
||||
PathCompletionController(PathCompleter(), self),
|
||||
])
|
||||
self._completion_popup: CompletionPopup | None = None
|
||||
self._body: ChatInputBody | None = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
self._completion_popup = CompletionPopup()
|
||||
yield self._completion_popup
|
||||
|
||||
with Vertical(
|
||||
id=self.ID_INPUT_BOX, classes="border-warning" if self._show_warning else ""
|
||||
):
|
||||
self._body = ChatInputBody(history_file=self._history_file, id="input-body")
|
||||
|
||||
yield self._body
|
||||
|
||||
def on_mount(self) -> None:
|
||||
if not self._body:
|
||||
return
|
||||
|
||||
self._body.set_completion_reset_callback(self._completion_manager.reset)
|
||||
if self._body.input_widget:
|
||||
self._body.input_widget.set_completion_manager(self._completion_manager)
|
||||
self._body.focus_input()
|
||||
|
||||
@property
|
||||
def input_widget(self) -> ChatTextArea | None:
|
||||
return self._body.input_widget if self._body else None
|
||||
|
||||
@property
|
||||
def value(self) -> str:
|
||||
if not self._body:
|
||||
return ""
|
||||
return self._body.value
|
||||
|
||||
@value.setter
|
||||
def value(self, text: str) -> None:
|
||||
if not self._body:
|
||||
return
|
||||
self._body.value = text
|
||||
widget = self._body.input_widget
|
||||
if widget:
|
||||
self._completion_manager.on_text_changed(
|
||||
widget.text, widget.get_cursor_offset()
|
||||
)
|
||||
|
||||
def focus_input(self) -> None:
|
||||
if self._body:
|
||||
self._body.focus_input()
|
||||
|
||||
def render_completion_suggestions(
|
||||
self, suggestions: list[tuple[str, str]], selected_index: int
|
||||
) -> None:
|
||||
if self._completion_popup:
|
||||
self._completion_popup.update_suggestions(suggestions, selected_index)
|
||||
|
||||
def clear_completion_suggestions(self) -> None:
|
||||
if self._completion_popup:
|
||||
self._completion_popup.hide()
|
||||
|
||||
def _format_insertion(self, replacement: str, suffix: str) -> str:
|
||||
"""Format the insertion text with appropriate spacing.
|
||||
|
||||
Args:
|
||||
replacement: The text to insert
|
||||
suffix: The text that follows the insertion point
|
||||
|
||||
Returns:
|
||||
The formatted insertion text with spacing if needed
|
||||
"""
|
||||
if replacement.startswith("@"):
|
||||
if replacement.endswith("/"):
|
||||
return replacement
|
||||
# For @-prefixed completions, add space unless suffix starts with whitespace
|
||||
return replacement + (" " if not suffix or not suffix[0].isspace() else "")
|
||||
|
||||
# For other completions, add space only if suffix exists and doesn't start with whitespace
|
||||
return replacement + (" " if suffix and not suffix[0].isspace() else "")
|
||||
|
||||
def replace_completion_range(self, start: int, end: int, replacement: str) -> None:
|
||||
widget = self.input_widget
|
||||
if not widget or not self._body:
|
||||
return
|
||||
|
||||
text = widget.text
|
||||
start = max(0, min(start, len(text)))
|
||||
end = max(start, min(end, len(text)))
|
||||
|
||||
prefix = text[:start]
|
||||
suffix = text[end:]
|
||||
insertion = self._format_insertion(replacement, suffix)
|
||||
new_text = f"{prefix}{insertion}{suffix}"
|
||||
|
||||
self._body.replace_input(new_text, cursor_offset=start + len(insertion))
|
||||
|
||||
def on_chat_input_body_submitted(self, event: ChatInputBody.Submitted) -> None:
|
||||
event.stop()
|
||||
self.post_message(self.Submitted(event.value))
|
||||
|
||||
def set_show_warning(self, show_warning: bool) -> None:
|
||||
self._show_warning = show_warning
|
||||
|
||||
input_box = self.get_widget_by_id(self.ID_INPUT_BOX)
|
||||
if show_warning:
|
||||
input_box.add_class(self.BORDER_WARNING_CLASS)
|
||||
else:
|
||||
input_box.remove_class(self.BORDER_WARNING_CLASS)
|
||||
246
vibe/cli/textual_ui/widgets/chat_input/text_area.py
Normal file
246
vibe/cli/textual_ui/widgets/chat_input/text_area.py
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from textual import events
|
||||
from textual.binding import Binding
|
||||
from textual.message import Message
|
||||
from textual.widgets import TextArea
|
||||
|
||||
from vibe.cli.autocompletion.base import CompletionResult
|
||||
from vibe.cli.textual_ui.widgets.chat_input.completion_manager import (
|
||||
MultiCompletionManager,
|
||||
)
|
||||
|
||||
|
||||
class ChatTextArea(TextArea):
|
||||
BINDINGS: ClassVar[list[Binding]] = [
|
||||
Binding(
|
||||
"shift+enter,ctrl+j",
|
||||
"insert_newline",
|
||||
"New Line",
|
||||
show=False,
|
||||
priority=True,
|
||||
)
|
||||
]
|
||||
|
||||
class Submitted(Message):
|
||||
def __init__(self, value: str) -> None:
|
||||
self.value = value
|
||||
super().__init__()
|
||||
|
||||
class HistoryPrevious(Message):
|
||||
def __init__(self, prefix: str) -> None:
|
||||
self.prefix = prefix
|
||||
super().__init__()
|
||||
|
||||
class HistoryNext(Message):
|
||||
def __init__(self, prefix: str) -> None:
|
||||
self.prefix = prefix
|
||||
super().__init__()
|
||||
|
||||
class HistoryReset(Message):
|
||||
"""Message sent when history navigation should be reset."""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self._history_prefix: str | None = None
|
||||
self._last_text = ""
|
||||
self._navigating_history = False
|
||||
self._last_cursor_col: int = 0
|
||||
self._last_used_prefix: str | None = None
|
||||
self._original_text: str = ""
|
||||
self._cursor_pos_after_load: tuple[int, int] | None = None
|
||||
self._cursor_moved_since_load: bool = False
|
||||
self._completion_manager: MultiCompletionManager | None = None
|
||||
|
||||
def on_blur(self, event: events.Blur) -> None:
|
||||
self.call_after_refresh(self.focus)
|
||||
|
||||
def on_click(self, event: events.Click) -> None:
|
||||
self._mark_cursor_moved_if_needed()
|
||||
|
||||
def action_insert_newline(self) -> None:
|
||||
self.insert("\n")
|
||||
|
||||
def on_text_area_changed(self, event: TextArea.Changed) -> None:
|
||||
if not self._navigating_history and self.text != self._last_text:
|
||||
self._reset_prefix()
|
||||
self._original_text = ""
|
||||
self._cursor_pos_after_load = None
|
||||
self._cursor_moved_since_load = False
|
||||
self.post_message(self.HistoryReset())
|
||||
self._last_text = self.text
|
||||
was_navigating_history = self._navigating_history
|
||||
self._navigating_history = False
|
||||
|
||||
if self._completion_manager and not was_navigating_history:
|
||||
self._completion_manager.on_text_changed(
|
||||
self.text, self.get_cursor_offset()
|
||||
)
|
||||
|
||||
def _reset_prefix(self) -> None:
|
||||
self._history_prefix = None
|
||||
self._last_used_prefix = None
|
||||
|
||||
def _mark_cursor_moved_if_needed(self) -> None:
|
||||
if (
|
||||
self._cursor_pos_after_load is not None
|
||||
and not self._cursor_moved_since_load
|
||||
and self.cursor_location != self._cursor_pos_after_load
|
||||
):
|
||||
self._cursor_moved_since_load = True
|
||||
self._reset_prefix()
|
||||
|
||||
def _get_prefix_up_to_cursor(self) -> str:
|
||||
cursor_row, cursor_col = self.cursor_location
|
||||
lines = self.text.split("\n")
|
||||
if cursor_row < len(lines):
|
||||
return lines[cursor_row][:cursor_col]
|
||||
return ""
|
||||
|
||||
def _handle_history_up(self) -> bool:
|
||||
cursor_row, cursor_col = self.cursor_location
|
||||
if cursor_row == 0:
|
||||
if self._history_prefix is not None and cursor_col != self._last_cursor_col:
|
||||
self._reset_prefix()
|
||||
self._last_cursor_col = 0
|
||||
|
||||
if self._history_prefix is None:
|
||||
self._history_prefix = self._get_prefix_up_to_cursor()
|
||||
|
||||
self._navigating_history = True
|
||||
self.post_message(self.HistoryPrevious(self._history_prefix))
|
||||
return True
|
||||
return False
|
||||
|
||||
def _handle_history_down(self) -> bool:
|
||||
cursor_row, cursor_col = self.cursor_location
|
||||
total_lines = self.text.count("\n") + 1
|
||||
|
||||
on_first_line_unmoved = cursor_row == 0 and not self._cursor_moved_since_load
|
||||
on_last_line = cursor_row == total_lines - 1
|
||||
|
||||
should_intercept = (
|
||||
on_first_line_unmoved and self._history_prefix is not None
|
||||
) or on_last_line
|
||||
|
||||
if not should_intercept:
|
||||
return False
|
||||
|
||||
if self._history_prefix is not None and cursor_col != self._last_cursor_col:
|
||||
self._reset_prefix()
|
||||
self._last_cursor_col = 0
|
||||
|
||||
if self._history_prefix is None:
|
||||
self._history_prefix = self._get_prefix_up_to_cursor()
|
||||
|
||||
self._navigating_history = True
|
||||
self.post_message(self.HistoryNext(self._history_prefix))
|
||||
return True
|
||||
|
||||
async def _on_key(self, event: events.Key) -> None:
|
||||
self._mark_cursor_moved_if_needed()
|
||||
|
||||
manager = self._completion_manager
|
||||
if manager:
|
||||
match manager.on_key(event, self.text, self.get_cursor_offset()):
|
||||
case CompletionResult.HANDLED:
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
return
|
||||
case CompletionResult.SUBMIT:
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
value = self.text.strip()
|
||||
if value:
|
||||
self._reset_prefix()
|
||||
self.post_message(self.Submitted(value))
|
||||
return
|
||||
|
||||
if event.key == "enter":
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
value = self.text.strip()
|
||||
if value:
|
||||
self._reset_prefix()
|
||||
self.post_message(self.Submitted(value))
|
||||
return
|
||||
|
||||
if event.key == "shift+enter":
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if event.key == "up" and self._handle_history_up():
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
if event.key == "down" and self._handle_history_down():
|
||||
event.prevent_default()
|
||||
event.stop()
|
||||
return
|
||||
|
||||
await super()._on_key(event)
|
||||
self._mark_cursor_moved_if_needed()
|
||||
|
||||
def set_completion_manager(self, manager: MultiCompletionManager | None) -> None:
|
||||
self._completion_manager = manager
|
||||
if self._completion_manager:
|
||||
self._completion_manager.on_text_changed(
|
||||
self.text, self.get_cursor_offset()
|
||||
)
|
||||
|
||||
def get_cursor_offset(self) -> int:
|
||||
text = self.text
|
||||
row, col = self.cursor_location
|
||||
|
||||
if not text:
|
||||
return 0
|
||||
|
||||
lines = text.split("\n")
|
||||
row = max(0, min(row, len(lines) - 1))
|
||||
col = max(0, col)
|
||||
|
||||
offset = sum(len(lines[i]) + 1 for i in range(row))
|
||||
return offset + min(col, len(lines[row]))
|
||||
|
||||
def set_cursor_offset(self, offset: int) -> None:
|
||||
text = self.text
|
||||
if offset <= 0:
|
||||
self.move_cursor((0, 0))
|
||||
return
|
||||
|
||||
if offset >= len(text):
|
||||
lines = text.split("\n")
|
||||
if not lines:
|
||||
self.move_cursor((0, 0))
|
||||
return
|
||||
last_row = len(lines) - 1
|
||||
self.move_cursor((last_row, len(lines[last_row])))
|
||||
return
|
||||
|
||||
remaining = offset
|
||||
lines = text.split("\n")
|
||||
|
||||
for row, line in enumerate(lines):
|
||||
line_length = len(line)
|
||||
if remaining <= line_length:
|
||||
self.move_cursor((row, remaining))
|
||||
return
|
||||
remaining -= line_length + 1
|
||||
|
||||
last_row = len(lines) - 1
|
||||
self.move_cursor((last_row, len(lines[last_row])))
|
||||
|
||||
def reset_history_state(self) -> None:
|
||||
self._reset_prefix()
|
||||
self._original_text = ""
|
||||
self._cursor_pos_after_load = None
|
||||
self._cursor_moved_since_load = False
|
||||
self._last_text = self.text
|
||||
|
||||
def clear_text(self) -> None:
|
||||
self.clear()
|
||||
self.reset_history_state()
|
||||
42
vibe/cli/textual_ui/widgets/compact.py
Normal file
42
vibe/cli/textual_ui/widgets/compact.py
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.textual_ui.widgets.blinking_message import BlinkingMessage
|
||||
|
||||
|
||||
class CompactMessage(BlinkingMessage):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.add_class("compact-message")
|
||||
self.old_tokens: int | None = None
|
||||
self.new_tokens: int | None = None
|
||||
self.error_message: str | None = None
|
||||
|
||||
def get_content(self) -> str:
|
||||
if self._is_blinking:
|
||||
return "Compacting conversation history..."
|
||||
|
||||
if self.error_message:
|
||||
return f"Error: {self.error_message}"
|
||||
|
||||
if self.old_tokens is not None and self.new_tokens is not None:
|
||||
reduction = self.old_tokens - self.new_tokens
|
||||
reduction_pct = (
|
||||
(reduction / self.old_tokens * 100) if self.old_tokens > 0 else 0
|
||||
)
|
||||
return (
|
||||
f"Compaction complete: {self.old_tokens:,} → "
|
||||
f"{self.new_tokens:,} tokens (-{reduction_pct:.1f}%)"
|
||||
)
|
||||
|
||||
return "Compaction complete"
|
||||
|
||||
def set_complete(
|
||||
self, old_tokens: int | None = None, new_tokens: int | None = None
|
||||
) -> None:
|
||||
self.old_tokens = old_tokens
|
||||
self.new_tokens = new_tokens
|
||||
self.stop_blinking(success=True)
|
||||
|
||||
def set_error(self, error_message: str) -> None:
|
||||
self.error_message = error_message
|
||||
self.stop_blinking(success=False)
|
||||
156
vibe/cli/textual_ui/widgets/config_app.py
Normal file
156
vibe/cli/textual_ui/widgets/config_app.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, ClassVar, TypedDict
|
||||
|
||||
from textual import events
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding, BindingType
|
||||
from textual.containers import Container, Vertical
|
||||
from textual.message import Message
|
||||
from textual.theme import BUILTIN_THEMES
|
||||
from textual.widgets import Static
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import VibeConfig
|
||||
|
||||
THEMES = sorted(k for k in BUILTIN_THEMES if k != "textual-ansi")
|
||||
|
||||
|
||||
class SettingDefinition(TypedDict):
|
||||
key: str
|
||||
label: str
|
||||
type: str
|
||||
options: list[str]
|
||||
value: str
|
||||
|
||||
|
||||
class ConfigApp(Container):
|
||||
can_focus = True
|
||||
can_focus_children = False
|
||||
|
||||
BINDINGS: ClassVar[list[BindingType]] = [
|
||||
Binding("up", "move_up", "Up", show=False),
|
||||
Binding("down", "move_down", "Down", show=False),
|
||||
Binding("space", "toggle_setting", "Toggle", show=False),
|
||||
Binding("enter", "cycle", "Next", show=False),
|
||||
]
|
||||
|
||||
class SettingChanged(Message):
|
||||
def __init__(self, key: str, value: str) -> None:
|
||||
super().__init__()
|
||||
self.key = key
|
||||
self.value = value
|
||||
|
||||
class ConfigClosed(Message):
|
||||
def __init__(self, changes: dict[str, str]) -> None:
|
||||
super().__init__()
|
||||
self.changes = changes
|
||||
|
||||
def __init__(self, config: VibeConfig) -> None:
|
||||
super().__init__(id="config-app")
|
||||
self.config = config
|
||||
self.selected_index = 0
|
||||
self.changes: dict[str, str] = {}
|
||||
|
||||
self.settings: list[SettingDefinition] = [
|
||||
{
|
||||
"key": "active_model",
|
||||
"label": "Model",
|
||||
"type": "cycle",
|
||||
"options": [m.alias for m in self.config.models],
|
||||
"value": self.config.active_model,
|
||||
},
|
||||
{
|
||||
"key": "textual_theme",
|
||||
"label": "Theme",
|
||||
"type": "cycle",
|
||||
"options": THEMES,
|
||||
"value": self.config.textual_theme,
|
||||
},
|
||||
]
|
||||
|
||||
self.title_widget: Static | None = None
|
||||
self.setting_widgets: list[Static] = []
|
||||
self.help_widget: Static | None = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Vertical(id="config-content"):
|
||||
self.title_widget = Static("Settings", classes="settings-title")
|
||||
yield self.title_widget
|
||||
|
||||
yield Static("")
|
||||
|
||||
for _ in self.settings:
|
||||
widget = Static("", classes="settings-option")
|
||||
self.setting_widgets.append(widget)
|
||||
yield widget
|
||||
|
||||
yield Static("")
|
||||
|
||||
self.help_widget = Static(
|
||||
"↑↓ navigate Space/Enter toggle ESC exit", classes="settings-help"
|
||||
)
|
||||
yield self.help_widget
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self._update_display()
|
||||
self.focus()
|
||||
|
||||
def _update_display(self) -> None:
|
||||
for i, (setting, widget) in enumerate(
|
||||
zip(self.settings, self.setting_widgets, strict=True)
|
||||
):
|
||||
is_selected = i == self.selected_index
|
||||
cursor = "› " if is_selected else " "
|
||||
|
||||
label: str = setting["label"]
|
||||
value: str = self.changes.get(setting["key"], setting["value"])
|
||||
|
||||
text = f"{cursor}{label}: {value}"
|
||||
|
||||
widget.update(text)
|
||||
|
||||
widget.remove_class("settings-cursor-selected")
|
||||
widget.remove_class("settings-value-cycle-selected")
|
||||
widget.remove_class("settings-value-cycle-unselected")
|
||||
|
||||
if is_selected:
|
||||
widget.add_class("settings-value-cycle-selected")
|
||||
else:
|
||||
widget.add_class("settings-value-cycle-unselected")
|
||||
|
||||
def action_move_up(self) -> None:
|
||||
self.selected_index = (self.selected_index - 1) % len(self.settings)
|
||||
self._update_display()
|
||||
|
||||
def action_move_down(self) -> None:
|
||||
self.selected_index = (self.selected_index + 1) % len(self.settings)
|
||||
self._update_display()
|
||||
|
||||
def action_toggle_setting(self) -> None:
|
||||
setting = self.settings[self.selected_index]
|
||||
key: str = setting["key"]
|
||||
current: str = self.changes.get(key, setting["value"])
|
||||
|
||||
options: list[str] = setting["options"]
|
||||
try:
|
||||
current_idx = options.index(current)
|
||||
next_idx = (current_idx + 1) % len(options)
|
||||
new_value: str = options[next_idx]
|
||||
except (ValueError, IndexError):
|
||||
new_value: str = options[0] if options else current
|
||||
|
||||
self.changes[key] = new_value
|
||||
|
||||
self.post_message(self.SettingChanged(key=key, value=new_value))
|
||||
|
||||
self._update_display()
|
||||
|
||||
def action_cycle(self) -> None:
|
||||
self.action_toggle_setting()
|
||||
|
||||
def action_close(self) -> None:
|
||||
self.post_message(self.ConfigClosed(changes=self.changes.copy()))
|
||||
|
||||
def on_blur(self, event: events.Blur) -> None:
|
||||
self.call_after_refresh(self.focus)
|
||||
31
vibe/cli/textual_ui/widgets/context_progress.py
Normal file
31
vibe/cli/textual_ui/widgets/context_progress.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from textual.reactive import reactive
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenState:
|
||||
max_tokens: int = 0
|
||||
current_tokens: int = 0
|
||||
|
||||
|
||||
class ContextProgress(Static):
|
||||
tokens = reactive(TokenState())
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def watch_tokens(self, new_state: TokenState) -> None:
|
||||
if new_state.max_tokens == 0:
|
||||
self.update("")
|
||||
return
|
||||
|
||||
percentage = min(
|
||||
100, int((new_state.current_tokens / new_state.max_tokens) * 100)
|
||||
)
|
||||
text = f"{percentage}% of {new_state.max_tokens // 1000}k tokens"
|
||||
self.update(text)
|
||||
157
vibe/cli/textual_ui/widgets/loading.py
Normal file
157
vibe/cli/textual_ui/widgets/loading.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
import random
|
||||
from time import time
|
||||
from typing import ClassVar
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
class LoadingWidget(Static):
|
||||
BRAILLE_SPINNER = ("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏")
|
||||
|
||||
TARGET_COLORS = ("#FFD800", "#FFAF00", "#FF8205", "#FA500F", "#E10500")
|
||||
|
||||
EASTER_EGGS: ClassVar[list[str]] = [
|
||||
"Eating a chocolatine",
|
||||
"Eating a pain au chocolat",
|
||||
"Réflexion",
|
||||
"Analyse",
|
||||
"Contemplation",
|
||||
"Synthèse",
|
||||
"Reading Proust",
|
||||
"Oui oui baguette",
|
||||
"Counting Rs in strawberry",
|
||||
"Seeding Mistral weights",
|
||||
"Vibing",
|
||||
"Sending good vibes",
|
||||
"Petting le chat",
|
||||
]
|
||||
|
||||
EASTER_EGGS_HALLOWEEN: ClassVar[list[str]] = [
|
||||
"Trick or treating",
|
||||
"Carving pumpkins",
|
||||
"Summoning spirits",
|
||||
"Brewing potions",
|
||||
"Haunting the terminal",
|
||||
"Petting le chat noir",
|
||||
]
|
||||
|
||||
EASTER_EGGS_DECEMBER: ClassVar[list[str]] = [
|
||||
"Wrapping presents",
|
||||
"Decorating the tree",
|
||||
"Drinking hot chocolate",
|
||||
"Building snowmen",
|
||||
"Writing holiday cards",
|
||||
]
|
||||
|
||||
def __init__(self, status: str | None = None) -> None:
|
||||
super().__init__(classes="loading-widget")
|
||||
self.status = status or self._get_default_status()
|
||||
self.gradient_offset = 0
|
||||
self.spinner_pos = 0
|
||||
self.char_widgets: list[Static] = []
|
||||
self.spinner_widget: Static | None = None
|
||||
self.ellipsis_widget: Static | None = None
|
||||
self.hint_widget: Static | None = None
|
||||
self.start_time: float | None = None
|
||||
|
||||
def _get_easter_egg(self) -> str | None:
|
||||
EASTER_EGG_PROBABILITY = 0.10
|
||||
if random.random() < EASTER_EGG_PROBABILITY:
|
||||
available_eggs = list(self.EASTER_EGGS)
|
||||
|
||||
OCTOBER = 10
|
||||
HALLOWEEN_DAY = 31
|
||||
DECEMBER = 12
|
||||
now = datetime.now()
|
||||
if now.month == OCTOBER and now.day == HALLOWEEN_DAY:
|
||||
available_eggs.extend(self.EASTER_EGGS_HALLOWEEN)
|
||||
if now.month == DECEMBER:
|
||||
available_eggs.extend(self.EASTER_EGGS_DECEMBER)
|
||||
|
||||
return random.choice(available_eggs)
|
||||
return None
|
||||
|
||||
def _get_default_status(self) -> str:
|
||||
return self._get_easter_egg() or "Thinking"
|
||||
|
||||
def _apply_easter_egg(self, status: str) -> str:
|
||||
return self._get_easter_egg() or status
|
||||
|
||||
def set_status(self, status: str) -> None:
|
||||
self.status = self._apply_easter_egg(status)
|
||||
self.gradient_offset = 0
|
||||
self._rebuild_chars()
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal(classes="loading-container"):
|
||||
self.spinner_widget = Static(
|
||||
self.BRAILLE_SPINNER[0] + " ", classes="loading-star"
|
||||
)
|
||||
yield self.spinner_widget
|
||||
|
||||
with Horizontal(classes="loading-status"):
|
||||
for char in self.status:
|
||||
widget = Static(char, classes="loading-char")
|
||||
self.char_widgets.append(widget)
|
||||
yield widget
|
||||
|
||||
self.ellipsis_widget = Static("… ", classes="loading-ellipsis")
|
||||
yield self.ellipsis_widget
|
||||
|
||||
self.hint_widget = Static("(0s esc to interrupt)", classes="loading-hint")
|
||||
yield self.hint_widget
|
||||
|
||||
def _rebuild_chars(self) -> None:
|
||||
if not self.is_mounted:
|
||||
return
|
||||
|
||||
status_container = self.query_one(".loading-status", Horizontal)
|
||||
|
||||
status_container.remove_children()
|
||||
self.char_widgets.clear()
|
||||
|
||||
for char in self.status:
|
||||
widget = Static(char, classes="loading-char")
|
||||
self.char_widgets.append(widget)
|
||||
status_container.mount(widget)
|
||||
|
||||
self.update_animation()
|
||||
|
||||
def on_mount(self) -> None:
|
||||
self.start_time = time()
|
||||
self.update_animation()
|
||||
self.set_interval(0.1, self.update_animation)
|
||||
|
||||
def _get_gradient_color(self, position: int) -> str:
|
||||
color_index = (position - self.gradient_offset) % len(self.TARGET_COLORS)
|
||||
return self.TARGET_COLORS[color_index]
|
||||
|
||||
def update_animation(self) -> None:
|
||||
if self.spinner_widget:
|
||||
spinner_char = self.BRAILLE_SPINNER[self.spinner_pos]
|
||||
color_0 = self._get_gradient_color(0)
|
||||
color_1 = self._get_gradient_color(1)
|
||||
self.spinner_widget.update(f"[{color_0}]{spinner_char}[/][{color_1}] [/]")
|
||||
self.spinner_pos = (self.spinner_pos + 1) % len(self.BRAILLE_SPINNER)
|
||||
|
||||
for i, widget in enumerate(self.char_widgets):
|
||||
position = 2 + i
|
||||
color = self._get_gradient_color(position)
|
||||
widget.update(f"[{color}]{self.status[i]}[/]")
|
||||
|
||||
if self.ellipsis_widget:
|
||||
ellipsis_start = 2 + len(self.status)
|
||||
color_ellipsis = self._get_gradient_color(ellipsis_start)
|
||||
color_space = self._get_gradient_color(ellipsis_start + 1)
|
||||
self.ellipsis_widget.update(f"[{color_ellipsis}]…[/][{color_space}] [/]")
|
||||
|
||||
self.gradient_offset = (self.gradient_offset + 1) % len(self.TARGET_COLORS)
|
||||
|
||||
if self.hint_widget and self.start_time is not None:
|
||||
elapsed = int(time() - self.start_time)
|
||||
self.hint_widget.update(f"({elapsed}s esc to interrupt)")
|
||||
148
vibe/cli/textual_ui/widgets/messages.py
Normal file
148
vibe/cli/textual_ui/widgets/messages.py
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.widgets import Markdown, Static
|
||||
from textual.widgets._markdown import MarkdownStream
|
||||
|
||||
|
||||
class UserMessage(Static):
|
||||
def __init__(self, content: str, pending: bool = False) -> None:
|
||||
super().__init__()
|
||||
self.add_class("user-message")
|
||||
self._content = content
|
||||
self._pending = pending
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal(classes="user-message-container"):
|
||||
yield Static("> ", classes="user-message-prompt")
|
||||
yield Static(self._content, markup=False, classes="user-message-content")
|
||||
if self._pending:
|
||||
self.add_class("pending")
|
||||
|
||||
async def set_pending(self, pending: bool) -> None:
|
||||
if pending == self._pending:
|
||||
return
|
||||
|
||||
self._pending = pending
|
||||
|
||||
if pending:
|
||||
self.add_class("pending")
|
||||
return
|
||||
|
||||
self.remove_class("pending")
|
||||
|
||||
|
||||
class AssistantMessage(Static):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__()
|
||||
self.add_class("assistant-message")
|
||||
self._content = content
|
||||
self._markdown: Markdown | None = None
|
||||
self._stream: MarkdownStream | None = None
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Horizontal(classes="assistant-message-container"):
|
||||
yield Static("● ", classes="assistant-message-dot")
|
||||
with Vertical(classes="assistant-message-content"):
|
||||
markdown = Markdown("")
|
||||
self._markdown = markdown
|
||||
yield markdown
|
||||
|
||||
def _get_markdown(self) -> Markdown:
|
||||
if self._markdown is None:
|
||||
self._markdown = self.query_one(Markdown)
|
||||
return self._markdown
|
||||
|
||||
def _ensure_stream(self) -> MarkdownStream:
|
||||
if self._stream is None:
|
||||
self._stream = Markdown.get_stream(self._get_markdown())
|
||||
return self._stream
|
||||
|
||||
async def append_content(self, content: str) -> None:
|
||||
if not content:
|
||||
return
|
||||
|
||||
self._content += content
|
||||
stream = self._ensure_stream()
|
||||
await stream.write(content)
|
||||
|
||||
async def write_initial_content(self) -> None:
|
||||
if self._content:
|
||||
stream = self._ensure_stream()
|
||||
await stream.write(self._content)
|
||||
|
||||
async def stop_stream(self) -> None:
|
||||
if self._stream is None:
|
||||
return
|
||||
|
||||
await self._stream.stop()
|
||||
self._stream = None
|
||||
|
||||
|
||||
class UserCommandMessage(Static):
|
||||
def __init__(self, content: str) -> None:
|
||||
super().__init__()
|
||||
self.add_class("user-command-message")
|
||||
self._content = content
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
yield Markdown(self._content)
|
||||
|
||||
|
||||
class InterruptMessage(Static):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
"Interrupted · What should Vibe do instead?", classes="interrupt-message"
|
||||
)
|
||||
|
||||
|
||||
class BashOutputMessage(Static):
|
||||
def __init__(self, command: str, cwd: str, output: str, exit_code: int) -> None:
|
||||
super().__init__()
|
||||
self.add_class("bash-output-message")
|
||||
self._command = command
|
||||
self._cwd = cwd
|
||||
self._output = output
|
||||
self._exit_code = exit_code
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Vertical(classes="bash-output-container"):
|
||||
with Horizontal(classes="bash-cwd-line"):
|
||||
yield Static(self._cwd, markup=False, classes="bash-cwd")
|
||||
yield Static("", classes="bash-cwd-spacer")
|
||||
if self._exit_code == 0:
|
||||
yield Static("✓", classes="bash-exit-success")
|
||||
else:
|
||||
yield Static("✗", classes="bash-exit-failure")
|
||||
yield Static(f" ({self._exit_code})", classes="bash-exit-code")
|
||||
with Horizontal(classes="bash-command-line"):
|
||||
yield Static("> ", classes="bash-chevron")
|
||||
yield Static(self._command, markup=False, classes="bash-command")
|
||||
yield Static("", classes="bash-command-spacer")
|
||||
yield Static(self._output, markup=False, classes="bash-output")
|
||||
|
||||
|
||||
class ErrorMessage(Static):
|
||||
def __init__(self, error: str, collapsed: bool = True) -> None:
|
||||
super().__init__(classes="error-message")
|
||||
self._error = error
|
||||
self.collapsed = collapsed
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
if self.collapsed:
|
||||
yield Static("Error. (ctrl+o to expand)", markup=False)
|
||||
else:
|
||||
yield Static(f"Error: {self._error}", markup=False)
|
||||
|
||||
def set_collapsed(self, collapsed: bool) -> None:
|
||||
if self.collapsed == collapsed:
|
||||
return
|
||||
|
||||
self.collapsed = collapsed
|
||||
self.remove_children()
|
||||
|
||||
if self.collapsed:
|
||||
self.mount(Static("Error. (ctrl+o to expand)", markup=False))
|
||||
else:
|
||||
self.mount(Static(f"Error: {self._error}", markup=False))
|
||||
25
vibe/cli/textual_ui/widgets/mode_indicator.py
Normal file
25
vibe/cli/textual_ui/widgets/mode_indicator.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
class ModeIndicator(Static):
|
||||
def __init__(self, auto_approve: bool = False) -> None:
|
||||
super().__init__()
|
||||
self.can_focus = False
|
||||
self._auto_approve = auto_approve
|
||||
self._update_display()
|
||||
|
||||
def _update_display(self) -> None:
|
||||
if self._auto_approve:
|
||||
self.update("⏵⏵ auto-approve on (shift+tab to toggle)")
|
||||
self.add_class("mode-on")
|
||||
self.remove_class("mode-off")
|
||||
else:
|
||||
self.update("⏵ auto-approve off (shift+tab to toggle)")
|
||||
self.add_class("mode-off")
|
||||
self.remove_class("mode-on")
|
||||
|
||||
def set_auto_approve(self, enabled: bool) -> None:
|
||||
self._auto_approve = enabled
|
||||
self._update_display()
|
||||
28
vibe/cli/textual_ui/widgets/path_display.py
Normal file
28
vibe/cli/textual_ui/widgets/path_display.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from textual.widgets import Static
|
||||
|
||||
|
||||
class PathDisplay(Static):
|
||||
def __init__(self, path: Path | str) -> None:
|
||||
super().__init__()
|
||||
self.can_focus = False
|
||||
self._path = Path(path)
|
||||
self._update_display()
|
||||
|
||||
def _update_display(self) -> None:
|
||||
path_str = str(self._path)
|
||||
try:
|
||||
home = Path.home()
|
||||
if self._path.is_relative_to(home):
|
||||
path_str = f"~/{self._path.relative_to(home)}"
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
|
||||
self.update(path_str)
|
||||
|
||||
def set_path(self, path: Path | str) -> None:
|
||||
self._path = Path(path)
|
||||
self._update_display()
|
||||
307
vibe/cli/textual_ui/widgets/tool_widgets.py
Normal file
307
vibe/cli/textual_ui/widgets/tool_widgets.py
Normal file
|
|
@ -0,0 +1,307 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.containers import Vertical
|
||||
from textual.widgets import Markdown, Static
|
||||
|
||||
|
||||
class ToolApprovalWidget(Vertical):
|
||||
def __init__(self, data: dict) -> None:
|
||||
super().__init__()
|
||||
self.data = data
|
||||
self.add_class("tool-approval-widget")
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
MAX_APPROVAL_MSG_SIZE = 150
|
||||
|
||||
for key, value in self.data.items():
|
||||
value_str = str(value)
|
||||
if len(value_str) > MAX_APPROVAL_MSG_SIZE:
|
||||
hidden = len(value_str) - MAX_APPROVAL_MSG_SIZE
|
||||
value_str = (
|
||||
value_str[:MAX_APPROVAL_MSG_SIZE] + f"… ({hidden} more characters)"
|
||||
)
|
||||
yield Static(
|
||||
f"{key}: {value_str}", markup=False, classes="approval-description"
|
||||
)
|
||||
|
||||
|
||||
class ToolResultWidget(Static):
|
||||
def __init__(self, data: dict, collapsed: bool = True) -> None:
|
||||
super().__init__()
|
||||
self.data = data
|
||||
self.collapsed = collapsed
|
||||
self.add_class("tool-result-widget")
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if not self.collapsed and (details := self.data.get("details")):
|
||||
for key, value in details.items():
|
||||
if value:
|
||||
yield Static(
|
||||
f"{key}: {value}", markup=False, classes="tool-result-detail"
|
||||
)
|
||||
|
||||
|
||||
class BashApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
command = self.data.get("command", "")
|
||||
description = self.data.get("description", "")
|
||||
|
||||
if description:
|
||||
yield Static(description, markup=False, classes="approval-description")
|
||||
yield Static("")
|
||||
|
||||
yield Markdown(f"```bash\n{command}\n```")
|
||||
|
||||
|
||||
class BashResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if not self.collapsed and (details := self.data.get("details")):
|
||||
for key, value in details.items():
|
||||
if value:
|
||||
yield Static(
|
||||
f"{key}: {value}", markup=False, classes="tool-result-detail"
|
||||
)
|
||||
|
||||
|
||||
class WriteFileApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
path = self.data.get("path", "")
|
||||
content = self.data.get("content", "")
|
||||
file_extension = self.data.get("file_extension", "text")
|
||||
|
||||
yield Static(f"File: {path}", markup=False, classes="approval-description")
|
||||
yield Static("")
|
||||
|
||||
yield Markdown(f"```{file_extension}\n{content}\n```")
|
||||
|
||||
|
||||
class WriteFileResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
MAX_LINES = 10
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if not self.collapsed:
|
||||
if path := self.data.get("path"):
|
||||
yield Static(
|
||||
f"Path: {path}", markup=False, classes="tool-result-detail"
|
||||
)
|
||||
|
||||
if bytes_written := self.data.get("bytes_written"):
|
||||
yield Static(
|
||||
f"Bytes: {bytes_written}",
|
||||
markup=False,
|
||||
classes="tool-result-detail",
|
||||
)
|
||||
|
||||
if content := self.data.get("content"):
|
||||
yield Static("")
|
||||
file_extension = self.data.get("file_extension", "text")
|
||||
|
||||
lines = content.split("\n")
|
||||
total_lines = len(lines)
|
||||
|
||||
if total_lines > MAX_LINES:
|
||||
shown_lines = lines[:MAX_LINES]
|
||||
remaining = total_lines - MAX_LINES
|
||||
truncated_content = "\n".join(
|
||||
shown_lines + [f"… ({remaining} more lines)"]
|
||||
)
|
||||
yield Markdown(f"```{file_extension}\n{truncated_content}\n```")
|
||||
else:
|
||||
yield Markdown(f"```{file_extension}\n{content}\n```")
|
||||
|
||||
|
||||
class SearchReplaceApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
file_path = self.data.get("file_path", "")
|
||||
diff_lines = self.data.get("diff_lines", [])
|
||||
|
||||
yield Static(f"File: {file_path}", markup=False, classes="approval-description")
|
||||
yield Static("")
|
||||
|
||||
if diff_lines:
|
||||
for line in diff_lines:
|
||||
if line.startswith("---") or line.startswith("+++"):
|
||||
yield Static(line, markup=False, classes="diff-header")
|
||||
elif line.startswith("-"):
|
||||
yield Static(line, markup=False, classes="diff-removed")
|
||||
elif line.startswith("+"):
|
||||
yield Static(line, markup=False, classes="diff-added")
|
||||
elif line.startswith("@@"):
|
||||
yield Static(line, markup=False, classes="diff-range")
|
||||
else:
|
||||
yield Static(line, markup=False, classes="diff-context")
|
||||
|
||||
|
||||
class SearchReplaceResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if not self.collapsed and (diff_lines := self.data.get("diff_lines")):
|
||||
yield Static("")
|
||||
for line in diff_lines:
|
||||
if line.startswith("---") or line.startswith("+++"):
|
||||
yield Static(line, markup=False, classes="diff-header")
|
||||
elif line.startswith("-"):
|
||||
yield Static(line, markup=False, classes="diff-removed")
|
||||
elif line.startswith("+"):
|
||||
yield Static(line, markup=False, classes="diff-added")
|
||||
elif line.startswith("@@"):
|
||||
yield Static(line, markup=False, classes="diff-range")
|
||||
else:
|
||||
yield Static(line, markup=False, classes="diff-context")
|
||||
|
||||
|
||||
class TodoApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
description = self.data.get("description", "")
|
||||
if description:
|
||||
yield Static(description, markup=False, classes="approval-description")
|
||||
|
||||
|
||||
class TodoResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(message, markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
yield Static("")
|
||||
|
||||
by_status = self.data.get("todos_by_status", {})
|
||||
if not any(by_status.values()):
|
||||
yield Static("No todos", markup=False, classes="todo-empty")
|
||||
return
|
||||
|
||||
for status in ["in_progress", "pending", "completed", "cancelled"]:
|
||||
todos = by_status.get(status, [])
|
||||
for todo in todos:
|
||||
content = todo.get("content", "")
|
||||
icon = self._get_status_icon(status)
|
||||
yield Static(
|
||||
f"{icon} {content}", markup=False, classes=f"todo-{status}"
|
||||
)
|
||||
|
||||
def _get_status_icon(self, status: str) -> str:
|
||||
icons = {"pending": "☐", "in_progress": "☐", "completed": "☑", "cancelled": "☒"}
|
||||
return icons.get(status, "☐")
|
||||
|
||||
|
||||
class ReadFileApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
for key, value in self.data.items():
|
||||
if value:
|
||||
yield Static(
|
||||
f"{key}: {value}", markup=False, classes="approval-description"
|
||||
)
|
||||
|
||||
|
||||
class ReadFileResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
MAX_LINES = 10
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if self.collapsed:
|
||||
return
|
||||
|
||||
if path := self.data.get("path"):
|
||||
yield Static(f"Path: {path}", markup=False, classes="tool-result-detail")
|
||||
|
||||
if warnings := self.data.get("warnings"):
|
||||
for warning in warnings:
|
||||
yield Static(
|
||||
f"⚠ {warning}", markup=False, classes="tool-result-warning"
|
||||
)
|
||||
|
||||
if content := self.data.get("content"):
|
||||
yield Static("")
|
||||
file_extension = self.data.get("file_extension", "text")
|
||||
|
||||
lines = content.split("\n")
|
||||
total_lines = len(lines)
|
||||
|
||||
if total_lines > MAX_LINES:
|
||||
shown_lines = lines[:MAX_LINES]
|
||||
remaining = total_lines - MAX_LINES
|
||||
truncated_content = "\n".join(
|
||||
shown_lines + [f"… ({remaining} more lines)"]
|
||||
)
|
||||
yield Markdown(f"```{file_extension}\n{truncated_content}\n```")
|
||||
else:
|
||||
yield Markdown(f"```{file_extension}\n{content}\n```")
|
||||
|
||||
|
||||
class GrepApprovalWidget(ToolApprovalWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
for key, value in self.data.items():
|
||||
if value:
|
||||
yield Static(
|
||||
f"{key}: {value!s}", classes="approval-description", markup=False
|
||||
)
|
||||
|
||||
|
||||
class GrepResultWidget(ToolResultWidget):
|
||||
def compose(self) -> ComposeResult:
|
||||
MAX_LINES = 30
|
||||
message = self.data.get("message", "")
|
||||
|
||||
if self.collapsed:
|
||||
yield Static(f"{message} (ctrl+o to expand.)", markup=False)
|
||||
else:
|
||||
yield Static(message, markup=False)
|
||||
|
||||
if self.collapsed:
|
||||
return
|
||||
|
||||
if warnings := self.data.get("warnings"):
|
||||
for warning in warnings:
|
||||
yield Static(
|
||||
f"⚠ {warning}", classes="tool-result-warning", markup=False
|
||||
)
|
||||
|
||||
if matches := self.data.get("matches"):
|
||||
yield Static("")
|
||||
|
||||
lines = matches.split("\n")
|
||||
total_lines = len(lines)
|
||||
|
||||
if total_lines > MAX_LINES:
|
||||
shown_lines = lines[:MAX_LINES]
|
||||
remaining = total_lines - MAX_LINES
|
||||
truncated_content = "\n".join(
|
||||
shown_lines + [f"… ({remaining} more lines)"]
|
||||
)
|
||||
yield Markdown(f"```\n{truncated_content}\n```")
|
||||
else:
|
||||
yield Markdown(f"```\n{matches}\n```")
|
||||
86
vibe/cli/textual_ui/widgets/tools.py
Normal file
86
vibe/cli/textual_ui/widgets/tools.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from textual.widgets import Static
|
||||
|
||||
from vibe.cli.textual_ui.renderers import get_renderer
|
||||
from vibe.cli.textual_ui.widgets.blinking_message import BlinkingMessage
|
||||
from vibe.core.tools.ui import ToolUIDataAdapter
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class ToolCallMessage(BlinkingMessage):
|
||||
def __init__(self, event: ToolCallEvent) -> None:
|
||||
self.event = event
|
||||
super().__init__()
|
||||
self.add_class("tool-call")
|
||||
|
||||
def get_content(self) -> str:
|
||||
if not self.event.tool_class:
|
||||
return f"{self.event.tool_name}"
|
||||
|
||||
adapter = ToolUIDataAdapter(self.event.tool_class)
|
||||
display = adapter.get_call_display(self.event)
|
||||
|
||||
return f"{display.summary}"
|
||||
|
||||
|
||||
class ToolResultMessage(Static):
|
||||
def __init__(
|
||||
self,
|
||||
event: ToolResultEvent,
|
||||
call_widget: ToolCallMessage | None = None,
|
||||
collapsed: bool = True,
|
||||
) -> None:
|
||||
self.event = event
|
||||
self.call_widget = call_widget
|
||||
self.collapsed = collapsed
|
||||
|
||||
super().__init__()
|
||||
self.add_class("tool-result")
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
if self.call_widget:
|
||||
success = not self.event.error and not self.event.skipped
|
||||
self.call_widget.stop_blinking(success=success)
|
||||
await self.render_result()
|
||||
|
||||
async def render_result(self) -> None:
|
||||
await self.remove_children()
|
||||
|
||||
if self.event.error:
|
||||
self.add_class("error-text")
|
||||
if self.collapsed:
|
||||
self.update("Error. (ctrl+o to expand)")
|
||||
else:
|
||||
await self.mount(Static(f"Error: {self.event.error}", markup=False))
|
||||
return
|
||||
|
||||
if self.event.skipped:
|
||||
self.add_class("warning-text")
|
||||
reason = self.event.skip_reason or "User skipped"
|
||||
if self.collapsed:
|
||||
self.update("Skipped. (ctrl+o to expand)")
|
||||
else:
|
||||
await self.mount(Static(f"Skipped: {reason}", markup=False))
|
||||
return
|
||||
|
||||
self.remove_class("error-text")
|
||||
self.remove_class("warning-text")
|
||||
|
||||
adapter = ToolUIDataAdapter(self.event.tool_class)
|
||||
display = adapter.get_result_display(self.event)
|
||||
|
||||
renderer = get_renderer(self.event.tool_name)
|
||||
widget_class, data = renderer.get_result_widget(display, self.collapsed)
|
||||
|
||||
result_widget = widget_class(data, collapsed=self.collapsed)
|
||||
await self.mount(result_widget)
|
||||
|
||||
async def set_collapsed(self, collapsed: bool) -> None:
|
||||
if self.collapsed != collapsed:
|
||||
self.collapsed = collapsed
|
||||
await self.render_result()
|
||||
|
||||
async def toggle_collapsed(self) -> None:
|
||||
self.collapsed = not self.collapsed
|
||||
await self.render_result()
|
||||
283
vibe/cli/textual_ui/widgets/welcome.py
Normal file
283
vibe/cli/textual_ui/widgets/welcome.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from time import monotonic
|
||||
|
||||
from rich.align import Align
|
||||
from rich.console import Group
|
||||
from rich.text import Text
|
||||
from textual.color import Color
|
||||
from textual.widgets import Static
|
||||
|
||||
from vibe.core import __version__
|
||||
from vibe.core.config import VibeConfig
|
||||
|
||||
|
||||
def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
|
||||
normalized = hex_color.lstrip("#")
|
||||
r, g, b = (int(normalized[i : i + 2], 16) for i in (0, 2, 4))
|
||||
return (r, g, b)
|
||||
|
||||
|
||||
def rgb_to_hex(r: int, g: int, b: int) -> str:
|
||||
return f"#{r:02x}{g:02x}{b:02x}"
|
||||
|
||||
|
||||
def interpolate_color(
|
||||
start_rgb: tuple[int, int, int], end_rgb: tuple[int, int, int], progress: float
|
||||
) -> str:
|
||||
progress = max(0.0, min(1.0, progress))
|
||||
r = int(start_rgb[0] + (end_rgb[0] - start_rgb[0]) * progress)
|
||||
g = int(start_rgb[1] + (end_rgb[1] - start_rgb[1]) * progress)
|
||||
b = int(start_rgb[2] + (end_rgb[2] - start_rgb[2]) * progress)
|
||||
return rgb_to_hex(r, g, b)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LineAnimationState:
|
||||
progress: float = 0.0
|
||||
cached_color: str | None = None
|
||||
cached_progress: float = -1.0
|
||||
rendered_color: str | None = None
|
||||
|
||||
|
||||
class WelcomeBanner(Static):
|
||||
FLASH_COLOR = "#FFFFFF"
|
||||
TARGET_COLORS = ("#FFD800", "#FFAF00", "#FF8205", "#FA500F", "#E10500")
|
||||
BORDER_TARGET_COLOR = "#b05800"
|
||||
|
||||
LINE_ANIMATION_DURATION_MS = 200
|
||||
LINE_STAGGER_MS = 280
|
||||
FLASH_RESET_DURATION_MS = 400
|
||||
ANIMATION_TICK_INTERVAL = 0.1
|
||||
|
||||
COLOR_FLASH_MIDPOINT = 0.5
|
||||
COLOR_PHASE_SCALE = 2.0
|
||||
COLOR_CACHE_THRESHOLD = 0.001
|
||||
BORDER_PROGRESS_THRESHOLD = 0.01
|
||||
|
||||
BLOCK = "▇▇"
|
||||
SPACE = " "
|
||||
LOGO_TEXT_GAP = " "
|
||||
|
||||
def __init__(self, config: VibeConfig) -> None:
|
||||
super().__init__(" ")
|
||||
self.config = config
|
||||
self.animation_timer = None
|
||||
self._animation_start_time: float | None = None
|
||||
|
||||
self._cached_skeleton_color: str | None = None
|
||||
self._cached_skeleton_rgb: tuple[int, int, int] | None = None
|
||||
self._flash_rgb = hex_to_rgb(self.FLASH_COLOR)
|
||||
self._target_rgbs = [hex_to_rgb(c) for c in self.TARGET_COLORS]
|
||||
self._border_target_rgb = hex_to_rgb(self.BORDER_TARGET_COLOR)
|
||||
|
||||
self._line_states = [LineAnimationState() for _ in self.TARGET_COLORS]
|
||||
self.border_progress = 0.0
|
||||
self._cached_border_color: str | None = None
|
||||
self._cached_border_progress = -1.0
|
||||
|
||||
self._line_duration = self.LINE_ANIMATION_DURATION_MS / 1000
|
||||
self._line_stagger = self.LINE_STAGGER_MS / 1000
|
||||
self._border_duration = self.FLASH_RESET_DURATION_MS / 1000
|
||||
self._line_start_times = [
|
||||
idx * self._line_stagger for idx in range(len(self.TARGET_COLORS))
|
||||
]
|
||||
self._all_lines_finish_time = (
|
||||
(len(self.TARGET_COLORS) - 1) * self.LINE_STAGGER_MS
|
||||
+ self.LINE_ANIMATION_DURATION_MS
|
||||
) / 1000
|
||||
|
||||
self._cached_text_lines: list[Text | None] = [None] * 7
|
||||
self._initialize_static_line_suffixes()
|
||||
|
||||
def _initialize_static_line_suffixes(self) -> None:
|
||||
self._static_line1_suffix = (
|
||||
f"{self.LOGO_TEXT_GAP}[b]Mistral Vibe v{__version__}[/]"
|
||||
)
|
||||
self._static_line2_suffix = (
|
||||
f"{self.LOGO_TEXT_GAP}[dim]{self.config.active_model}[/]"
|
||||
)
|
||||
mcp_count = len(self.config.mcp_servers)
|
||||
model_count = len(self.config.models)
|
||||
self._static_line3_suffix = f"{self.LOGO_TEXT_GAP}[dim]{model_count} models · {mcp_count} MCP servers[/]"
|
||||
self._static_line5_suffix = (
|
||||
f"{self.LOGO_TEXT_GAP}[dim]{self.config.effective_workdir}[/]"
|
||||
)
|
||||
block = (self.SPACE * 4) + self.LOGO_TEXT_GAP
|
||||
self._static_line7 = f"{block}[dim]Type[/] [{self.BORDER_TARGET_COLOR}]/help[/] [dim]for more information[/]"
|
||||
|
||||
@property
|
||||
def skeleton_color(self) -> str:
|
||||
return self._cached_skeleton_color or "#1e1e1e"
|
||||
|
||||
@property
|
||||
def skeleton_rgb(self) -> tuple[int, int, int]:
|
||||
return self._cached_skeleton_rgb or hex_to_rgb("#1e1e1e")
|
||||
|
||||
def on_mount(self) -> None:
|
||||
if not self.config.disable_welcome_banner_animation:
|
||||
self.call_after_refresh(self._init_after_styles)
|
||||
|
||||
def _init_after_styles(self) -> None:
|
||||
self._cache_skeleton_color()
|
||||
self._cached_text_lines[5] = Text("")
|
||||
self._cached_text_lines[6] = Text.from_markup(self._static_line7)
|
||||
self._update_display()
|
||||
self._start_animation()
|
||||
|
||||
def _cache_skeleton_color(self) -> None:
|
||||
try:
|
||||
border = self.styles.border
|
||||
if (
|
||||
hasattr(border, "top")
|
||||
and isinstance(edge := border.top, tuple)
|
||||
and len(edge) >= 2 # noqa: PLR2004
|
||||
and isinstance(color := edge[1], Color)
|
||||
):
|
||||
self._cached_skeleton_color = color.hex
|
||||
self._cached_skeleton_rgb = hex_to_rgb(color.hex)
|
||||
return
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
self._cached_skeleton_color = "#1e1e1e"
|
||||
self._cached_skeleton_rgb = hex_to_rgb("#1e1e1e")
|
||||
|
||||
def _stop_timer(self) -> None:
|
||||
if self.animation_timer:
|
||||
try:
|
||||
self.animation_timer.stop()
|
||||
except Exception:
|
||||
pass
|
||||
self.animation_timer = None
|
||||
|
||||
def on_unmount(self) -> None:
|
||||
self._stop_timer()
|
||||
|
||||
def _start_animation(self) -> None:
|
||||
self._animation_start_time = monotonic()
|
||||
|
||||
def tick() -> None:
|
||||
if self._is_animation_complete():
|
||||
self._stop_timer()
|
||||
return
|
||||
if self._animation_start_time is None:
|
||||
return
|
||||
|
||||
elapsed = monotonic() - self._animation_start_time
|
||||
updated_lines = self._advance_line_progress(elapsed)
|
||||
border_updated = self._advance_border_progress(elapsed)
|
||||
|
||||
if border_updated:
|
||||
self._update_border_color()
|
||||
if updated_lines or border_updated:
|
||||
self._update_display()
|
||||
|
||||
self.animation_timer = self.set_interval(self.ANIMATION_TICK_INTERVAL, tick)
|
||||
|
||||
def _advance_line_progress(self, elapsed: float) -> bool:
|
||||
any_updates = False
|
||||
for line_idx, state in enumerate(self._line_states):
|
||||
if state.progress >= 1.0:
|
||||
continue
|
||||
start_time = self._line_start_times[line_idx]
|
||||
if elapsed < start_time:
|
||||
continue
|
||||
progress = min(1.0, (elapsed - start_time) / self._line_duration)
|
||||
if progress > state.progress:
|
||||
state.progress = progress
|
||||
any_updates = True
|
||||
return any_updates
|
||||
|
||||
def _advance_border_progress(self, elapsed: float) -> bool:
|
||||
if elapsed < self._all_lines_finish_time:
|
||||
return False
|
||||
|
||||
new_progress = min(
|
||||
1.0, (elapsed - self._all_lines_finish_time) / self._border_duration
|
||||
)
|
||||
|
||||
if abs(new_progress - self.border_progress) > self.BORDER_PROGRESS_THRESHOLD:
|
||||
self.border_progress = new_progress
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _is_animation_complete(self) -> bool:
|
||||
return (
|
||||
all(state.progress >= 1.0 for state in self._line_states)
|
||||
and self.border_progress >= 1.0
|
||||
)
|
||||
|
||||
def _update_border_color(self) -> None:
|
||||
progress = self.border_progress
|
||||
if abs(progress - self._cached_border_progress) < self.COLOR_CACHE_THRESHOLD:
|
||||
return
|
||||
|
||||
border_color = self._compute_color_for_progress(
|
||||
progress, self._border_target_rgb
|
||||
)
|
||||
self._cached_border_color = border_color
|
||||
self._cached_border_progress = progress
|
||||
self.styles.border = ("round", border_color)
|
||||
|
||||
def _compute_color_for_progress(
|
||||
self, progress: float, target_rgb: tuple[int, int, int]
|
||||
) -> str:
|
||||
if progress <= 0:
|
||||
return self.skeleton_color
|
||||
|
||||
if progress <= self.COLOR_FLASH_MIDPOINT:
|
||||
phase = progress * self.COLOR_PHASE_SCALE
|
||||
return interpolate_color(self.skeleton_rgb, self._flash_rgb, phase)
|
||||
|
||||
phase = (progress - self.COLOR_FLASH_MIDPOINT) * self.COLOR_PHASE_SCALE
|
||||
return interpolate_color(self._flash_rgb, target_rgb, phase)
|
||||
|
||||
def _update_display(self) -> None:
|
||||
for idx in range(5):
|
||||
self._update_colored_line(idx, idx)
|
||||
|
||||
lines = [line if line else Text("") for line in self._cached_text_lines]
|
||||
self.update(Align.center(Group(*lines)))
|
||||
|
||||
def _get_color(self, line_idx: int) -> str:
|
||||
state = self._line_states[line_idx]
|
||||
if (
|
||||
abs(state.progress - state.cached_progress) < self.COLOR_CACHE_THRESHOLD
|
||||
and state.cached_color
|
||||
):
|
||||
return state.cached_color
|
||||
|
||||
color = self._compute_color_for_progress(
|
||||
state.progress, self._target_rgbs[line_idx]
|
||||
)
|
||||
state.cached_color = color
|
||||
state.cached_progress = state.progress
|
||||
return color
|
||||
|
||||
def _update_colored_line(self, slot_idx: int, line_idx: int) -> None:
|
||||
color = self._get_color(line_idx)
|
||||
state = self._line_states[line_idx]
|
||||
|
||||
if color == state.rendered_color and self._cached_text_lines[slot_idx]:
|
||||
return
|
||||
|
||||
state.rendered_color = color
|
||||
self._cached_text_lines[slot_idx] = Text.from_markup(
|
||||
self._build_line(slot_idx, color)
|
||||
)
|
||||
|
||||
def _build_line(self, line_idx: int, color: str) -> str:
|
||||
B = self.BLOCK
|
||||
S = self.SPACE
|
||||
|
||||
patterns = [
|
||||
f"{S}[{color}]{B}[/]{S}{S}{S}[{color}]{B}[/]{S}{self._static_line1_suffix}",
|
||||
f"{S}[{color}]{B}{B}[/]{S}[{color}]{B}{B}[/]{S}{self._static_line2_suffix}",
|
||||
f"{S}[{color}]{B}{B}{B}{B}{B}[/]{S}{self._static_line3_suffix}",
|
||||
f"{S}[{color}]{B}[/]{S}[{color}]{B}[/]{S}[{color}]{B}[/]{S}",
|
||||
f"[{color}]{B}{B}{B}[/]{S}[{color}]{B}{B}{B}[/]{self._static_line5_suffix}",
|
||||
]
|
||||
return patterns[line_idx]
|
||||
31
vibe/cli/update_notifier/__init__.py
Normal file
31
vibe/cli/update_notifier/__init__.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.update_notifier.fake_version_update_gateway import (
|
||||
FakeVersionUpdateGateway,
|
||||
)
|
||||
from vibe.cli.update_notifier.github_version_update_gateway import (
|
||||
GitHubVersionUpdateGateway,
|
||||
)
|
||||
from vibe.cli.update_notifier.version_update import (
|
||||
VersionUpdateError,
|
||||
is_version_update_available,
|
||||
)
|
||||
from vibe.cli.update_notifier.version_update_gateway import (
|
||||
DEFAULT_GATEWAY_MESSAGES,
|
||||
VersionUpdate,
|
||||
VersionUpdateGateway,
|
||||
VersionUpdateGatewayCause,
|
||||
VersionUpdateGatewayError,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_GATEWAY_MESSAGES",
|
||||
"FakeVersionUpdateGateway",
|
||||
"GitHubVersionUpdateGateway",
|
||||
"VersionUpdate",
|
||||
"VersionUpdateError",
|
||||
"VersionUpdateGateway",
|
||||
"VersionUpdateGatewayCause",
|
||||
"VersionUpdateGatewayError",
|
||||
"is_version_update_available",
|
||||
]
|
||||
24
vibe/cli/update_notifier/fake_version_update_gateway.py
Normal file
24
vibe/cli/update_notifier/fake_version_update_gateway.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.cli.update_notifier.version_update_gateway import (
|
||||
VersionUpdate,
|
||||
VersionUpdateGateway,
|
||||
VersionUpdateGatewayError,
|
||||
)
|
||||
|
||||
|
||||
class FakeVersionUpdateGateway(VersionUpdateGateway):
|
||||
def __init__(
|
||||
self,
|
||||
update: VersionUpdate | None = None,
|
||||
error: VersionUpdateGatewayError | None = None,
|
||||
) -> None:
|
||||
self._update: VersionUpdate | None = update
|
||||
self._error = error
|
||||
self.fetch_update_calls = 0
|
||||
|
||||
async def fetch_update(self) -> VersionUpdate | None:
|
||||
self.fetch_update_calls += 1
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._update
|
||||
109
vibe/cli/update_notifier/github_version_update_gateway.py
Normal file
109
vibe/cli/update_notifier/github_version_update_gateway.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
from vibe.cli.update_notifier.version_update_gateway import (
|
||||
VersionUpdate,
|
||||
VersionUpdateGateway,
|
||||
VersionUpdateGatewayCause,
|
||||
VersionUpdateGatewayError,
|
||||
)
|
||||
|
||||
|
||||
class GitHubVersionUpdateGateway(VersionUpdateGateway):
|
||||
def __init__(
|
||||
self,
|
||||
owner: str,
|
||||
repository: str,
|
||||
*,
|
||||
token: str | None = None,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 5.0,
|
||||
base_url: str = "https://api.github.com",
|
||||
) -> None:
|
||||
self._owner = owner
|
||||
self._repository = repository
|
||||
self._token = token
|
||||
self._client = client
|
||||
self._timeout = timeout
|
||||
self._base_url = base_url.rstrip("/")
|
||||
|
||||
async def fetch_update(self) -> VersionUpdate | None:
|
||||
headers = {
|
||||
"Accept": "application/vnd.github+json",
|
||||
"User-Agent": "mistral-vibe-update-notifier",
|
||||
}
|
||||
if self._token:
|
||||
headers["Authorization"] = f"Bearer {self._token}"
|
||||
|
||||
request_path = f"/repos/{self._owner}/{self._repository}/releases"
|
||||
|
||||
try:
|
||||
if self._client is not None:
|
||||
response = await self._client.get(
|
||||
f"{self._base_url}{request_path}",
|
||||
headers=headers,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
else:
|
||||
async with httpx.AsyncClient(
|
||||
base_url=self._base_url, timeout=self._timeout
|
||||
) as client:
|
||||
response = await client.get(request_path, headers=headers)
|
||||
except httpx.RequestError as exc:
|
||||
raise VersionUpdateGatewayError(
|
||||
cause=VersionUpdateGatewayCause.REQUEST_FAILED
|
||||
) from exc
|
||||
|
||||
rate_limit_remaining = response.headers.get("X-RateLimit-Remaining")
|
||||
if response.status_code == httpx.codes.TOO_MANY_REQUESTS or (
|
||||
rate_limit_remaining is not None and rate_limit_remaining == "0"
|
||||
):
|
||||
raise VersionUpdateGatewayError(
|
||||
cause=VersionUpdateGatewayCause.TOO_MANY_REQUESTS
|
||||
)
|
||||
|
||||
if response.status_code == httpx.codes.FORBIDDEN:
|
||||
raise VersionUpdateGatewayError(cause=VersionUpdateGatewayCause.FORBIDDEN)
|
||||
|
||||
if response.status_code == httpx.codes.NOT_FOUND:
|
||||
raise VersionUpdateGatewayError(
|
||||
cause=VersionUpdateGatewayCause.NOT_FOUND,
|
||||
message="Unable to fetch the GitHub releases. Did you export a GITHUB_TOKEN environment variable?",
|
||||
)
|
||||
|
||||
if response.is_error:
|
||||
raise VersionUpdateGatewayError(
|
||||
cause=VersionUpdateGatewayCause.ERROR_RESPONSE
|
||||
)
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except ValueError as exc:
|
||||
raise VersionUpdateGatewayError(
|
||||
cause=VersionUpdateGatewayCause.INVALID_RESPONSE
|
||||
) from exc
|
||||
|
||||
if not data:
|
||||
return None
|
||||
|
||||
# pick the most recently published non-prerelease and non-draft release
|
||||
# github "list releases" API most likely returns ordered results, but this is not guaranteed
|
||||
for release in sorted(
|
||||
data, key=lambda x: x.get("published_at") or "", reverse=True
|
||||
):
|
||||
if release.get("prerelease") or release.get("draft"):
|
||||
continue
|
||||
if version := _extract_version(release.get("tag_name")):
|
||||
return VersionUpdate(latest_version=version)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _extract_version(tag_name: str | None) -> str | None:
|
||||
if not tag_name:
|
||||
return None
|
||||
tag = tag_name.strip()
|
||||
if not tag:
|
||||
return None
|
||||
return tag[1:] if tag.startswith(("v", "V")) else tag
|
||||
57
vibe/cli/update_notifier/version_update.py
Normal file
57
vibe/cli/update_notifier/version_update.py
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from packaging.version import InvalidVersion, Version
|
||||
|
||||
from vibe.cli.update_notifier.version_update_gateway import (
|
||||
DEFAULT_GATEWAY_MESSAGES,
|
||||
VersionUpdate,
|
||||
VersionUpdateGateway,
|
||||
VersionUpdateGatewayCause,
|
||||
VersionUpdateGatewayError,
|
||||
)
|
||||
|
||||
|
||||
class VersionUpdateError(Exception):
|
||||
def __init__(self, message: str) -> None:
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
def _parse_version(raw: str) -> Version | None:
|
||||
try:
|
||||
return Version(raw.replace("-", "+"))
|
||||
except InvalidVersion:
|
||||
return None
|
||||
|
||||
|
||||
def _describe_gateway_error(error: VersionUpdateGatewayError) -> str:
|
||||
if message := getattr(error, "user_message", None):
|
||||
return message
|
||||
|
||||
cause = getattr(error, "cause", VersionUpdateGatewayCause.UNKNOWN)
|
||||
if isinstance(cause, VersionUpdateGatewayCause):
|
||||
return DEFAULT_GATEWAY_MESSAGES.get(
|
||||
cause, DEFAULT_GATEWAY_MESSAGES[VersionUpdateGatewayCause.UNKNOWN]
|
||||
)
|
||||
|
||||
return DEFAULT_GATEWAY_MESSAGES[VersionUpdateGatewayCause.UNKNOWN]
|
||||
|
||||
|
||||
async def is_version_update_available(
|
||||
version_update_notifier: VersionUpdateGateway, current_version: str
|
||||
) -> VersionUpdate | None:
|
||||
try:
|
||||
update = await version_update_notifier.fetch_update()
|
||||
except VersionUpdateGatewayError as error:
|
||||
raise VersionUpdateError(_describe_gateway_error(error)) from error
|
||||
|
||||
if not update:
|
||||
return None
|
||||
|
||||
latest_version = _parse_version(update.latest_version)
|
||||
current = _parse_version(current_version)
|
||||
|
||||
if latest_version is None or current is None:
|
||||
return None
|
||||
|
||||
return update if latest_version > current else None
|
||||
54
vibe/cli/update_notifier/version_update_gateway.py
Normal file
54
vibe/cli/update_notifier/version_update_gateway.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum, auto
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VersionUpdate:
|
||||
latest_version: str
|
||||
|
||||
|
||||
class VersionUpdateGatewayCause(StrEnum):
|
||||
@staticmethod
|
||||
def _generate_next_value_(
|
||||
name: str, start: int, count: int, last_values: list[str]
|
||||
) -> str:
|
||||
return name.lower()
|
||||
|
||||
TOO_MANY_REQUESTS = auto()
|
||||
FORBIDDEN = auto()
|
||||
NOT_FOUND = auto()
|
||||
REQUEST_FAILED = auto()
|
||||
ERROR_RESPONSE = auto()
|
||||
INVALID_RESPONSE = auto()
|
||||
UNKNOWN = auto()
|
||||
|
||||
|
||||
DEFAULT_GATEWAY_MESSAGES: dict[VersionUpdateGatewayCause, str] = {
|
||||
VersionUpdateGatewayCause.TOO_MANY_REQUESTS: "Rate limit exceeded while checking for updates.",
|
||||
VersionUpdateGatewayCause.FORBIDDEN: "Request was forbidden while checking for updates.",
|
||||
VersionUpdateGatewayCause.NOT_FOUND: "Unable to fetch the releases. Please check your permissions.",
|
||||
VersionUpdateGatewayCause.REQUEST_FAILED: "Network error while checking for updates.",
|
||||
VersionUpdateGatewayCause.ERROR_RESPONSE: "Unexpected response received while checking for updates.",
|
||||
VersionUpdateGatewayCause.INVALID_RESPONSE: "Received an invalid response while checking for updates.",
|
||||
VersionUpdateGatewayCause.UNKNOWN: "Unable to determine whether an update is available.",
|
||||
}
|
||||
|
||||
|
||||
class VersionUpdateGatewayError(Exception):
|
||||
def __init__(
|
||||
self, *, cause: VersionUpdateGatewayCause, message: str | None = None
|
||||
) -> None:
|
||||
self.cause = cause
|
||||
self.user_message = message
|
||||
detail = message or DEFAULT_GATEWAY_MESSAGES.get(
|
||||
cause, DEFAULT_GATEWAY_MESSAGES[VersionUpdateGatewayCause.UNKNOWN]
|
||||
)
|
||||
super().__init__(detail)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VersionUpdateGateway(Protocol):
|
||||
async def fetch_update(self) -> VersionUpdate | None: ...
|
||||
6
vibe/core/__init__.py
Normal file
6
vibe/core/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
__all__ = ["__version__", "run_programmatic"]
|
||||
__version__ = "0.1.0"
|
||||
|
||||
from vibe.core.programmatic import run_programmatic
|
||||
984
vibe/core/agent.py
Normal file
984
vibe/core/agent.py
Normal file
|
|
@ -0,0 +1,984 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from enum import StrEnum, auto
|
||||
import time
|
||||
from typing import Any, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from vibe.core.config import VibeConfig
|
||||
from vibe.core.interaction_logger import InteractionLogger
|
||||
from vibe.core.llm.backend.factory import BACKEND_FACTORY
|
||||
from vibe.core.llm.format import APIToolFormatHandler, ResolvedMessage
|
||||
from vibe.core.llm.types import BackendLike
|
||||
from vibe.core.middleware import (
|
||||
AutoCompactMiddleware,
|
||||
ContextWarningMiddleware,
|
||||
ConversationContext,
|
||||
MiddlewareAction,
|
||||
MiddlewarePipeline,
|
||||
MiddlewareResult,
|
||||
PriceLimitMiddleware,
|
||||
ResetReason,
|
||||
TurnLimitMiddleware,
|
||||
)
|
||||
from vibe.core.prompts import UtilityPrompt
|
||||
from vibe.core.system_prompt import get_universal_system_prompt
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
ToolPermissionError,
|
||||
)
|
||||
from vibe.core.tools.manager import ToolManager
|
||||
from vibe.core.types import (
|
||||
AgentStats,
|
||||
ApprovalCallback,
|
||||
AssistantEvent,
|
||||
BaseEvent,
|
||||
CompactEndEvent,
|
||||
CompactStartEvent,
|
||||
LLMChunk,
|
||||
LLMMessage,
|
||||
Role,
|
||||
SyncApprovalCallback,
|
||||
ToolCall,
|
||||
ToolCallEvent,
|
||||
ToolResultEvent,
|
||||
)
|
||||
from vibe.core.utils import (
|
||||
TOOL_ERROR_TAG,
|
||||
VIBE_STOP_EVENT_TAG,
|
||||
ApprovalResponse,
|
||||
CancellationReason,
|
||||
get_user_agent,
|
||||
get_user_cancellation_message,
|
||||
is_user_cancellation_event,
|
||||
)
|
||||
|
||||
|
||||
class ToolExecutionResponse(StrEnum):
|
||||
SKIP = auto()
|
||||
EXECUTE = auto()
|
||||
|
||||
|
||||
class ToolDecision(BaseModel):
|
||||
verdict: ToolExecutionResponse
|
||||
feedback: str | None = None
|
||||
|
||||
|
||||
class AgentError(Exception):
|
||||
"""Base exception for Agent errors."""
|
||||
|
||||
|
||||
class AgentStateError(AgentError):
|
||||
"""Raised when agent is in an invalid state."""
|
||||
|
||||
|
||||
class LLMResponseError(AgentError):
|
||||
"""Raised when LLM response is malformed or missing expected data."""
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(
|
||||
self,
|
||||
config: VibeConfig,
|
||||
auto_approve: bool = False,
|
||||
message_observer: Callable[[LLMMessage], None] | None = None,
|
||||
max_turns: int | None = None,
|
||||
max_price: float | None = None,
|
||||
backend: BackendLike | None = None,
|
||||
enable_streaming: bool = False,
|
||||
) -> None:
|
||||
self.config = config
|
||||
|
||||
self.tool_manager = ToolManager(config)
|
||||
self.format_handler = APIToolFormatHandler()
|
||||
|
||||
self.backend_factory = lambda: backend or self._select_backend()
|
||||
self.backend = self.backend_factory()
|
||||
|
||||
self.message_observer = message_observer
|
||||
self._last_observed_message_index: int = 0
|
||||
self.middleware_pipeline = MiddlewarePipeline()
|
||||
self.enable_streaming = enable_streaming
|
||||
self._setup_middleware(max_turns, max_price)
|
||||
|
||||
system_prompt = get_universal_system_prompt(self.tool_manager, config)
|
||||
|
||||
self.messages = [LLMMessage(role=Role.system, content=system_prompt)]
|
||||
|
||||
if self.message_observer:
|
||||
self.message_observer(self.messages[0])
|
||||
self._last_observed_message_index = 1
|
||||
|
||||
self.stats = AgentStats()
|
||||
try:
|
||||
active_model = config.get_active_model()
|
||||
self.stats.input_price_per_million = active_model.input_price
|
||||
self.stats.output_price_per_million = active_model.output_price
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
self.auto_approve = auto_approve
|
||||
self.approval_callback: ApprovalCallback | None = None
|
||||
|
||||
self.session_id = str(uuid4())
|
||||
|
||||
self.interaction_logger = InteractionLogger(
|
||||
config.session_logging,
|
||||
self.session_id,
|
||||
auto_approve,
|
||||
config.effective_workdir,
|
||||
)
|
||||
|
||||
self._last_chunk: LLMChunk | None = None
|
||||
|
||||
def _select_backend(self) -> BackendLike:
|
||||
active_model = self.config.get_active_model()
|
||||
provider = self.config.get_provider_for_model(active_model)
|
||||
timeout = self.config.api_timeout
|
||||
return BACKEND_FACTORY[provider.backend](provider=provider, timeout=timeout)
|
||||
|
||||
def add_message(self, message: LLMMessage) -> None:
|
||||
self.messages.append(message)
|
||||
|
||||
def _flush_new_messages(self) -> None:
|
||||
if not self.message_observer:
|
||||
return
|
||||
|
||||
if self._last_observed_message_index >= len(self.messages):
|
||||
return
|
||||
|
||||
for msg in self.messages[self._last_observed_message_index :]:
|
||||
self.message_observer(msg)
|
||||
self._last_observed_message_index = len(self.messages)
|
||||
|
||||
async def act(self, msg: str) -> AsyncGenerator[BaseEvent]:
|
||||
self._clean_message_history()
|
||||
async for event in self._conversation_loop(msg):
|
||||
yield event
|
||||
|
||||
def _setup_middleware(self, max_turns: int | None, max_price: float | None) -> None:
|
||||
self.middleware_pipeline.clear()
|
||||
|
||||
if max_turns is not None:
|
||||
self.middleware_pipeline.add(TurnLimitMiddleware(max_turns))
|
||||
|
||||
if max_price is not None:
|
||||
self.middleware_pipeline.add(PriceLimitMiddleware(max_price))
|
||||
|
||||
if self.config.auto_compact_threshold > 0:
|
||||
self.middleware_pipeline.add(
|
||||
AutoCompactMiddleware(self.config.auto_compact_threshold)
|
||||
)
|
||||
if self.config.context_warnings:
|
||||
self.middleware_pipeline.add(
|
||||
ContextWarningMiddleware(0.5, self.config.auto_compact_threshold)
|
||||
)
|
||||
|
||||
async def _handle_middleware_result(
|
||||
self, result: MiddlewareResult
|
||||
) -> AsyncGenerator[BaseEvent]:
|
||||
match result.action:
|
||||
case MiddlewareAction.STOP:
|
||||
yield AssistantEvent(
|
||||
content=f"<{VIBE_STOP_EVENT_TAG}>{result.reason}</{VIBE_STOP_EVENT_TAG}>",
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
session_total_tokens=self.stats.session_total_llm_tokens,
|
||||
last_turn_duration=0,
|
||||
tokens_per_second=0,
|
||||
stopped_by_middleware=True,
|
||||
)
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
case MiddlewareAction.INJECT_MESSAGE:
|
||||
if result.message and len(self.messages) > 0:
|
||||
last_msg = self.messages[-1]
|
||||
if last_msg.content:
|
||||
last_msg.content += f"\n\n{result.message}"
|
||||
else:
|
||||
last_msg.content = result.message
|
||||
|
||||
case MiddlewareAction.COMPACT:
|
||||
old_tokens = result.metadata.get(
|
||||
"old_tokens", self.stats.context_tokens
|
||||
)
|
||||
threshold = result.metadata.get(
|
||||
"threshold", self.config.auto_compact_threshold
|
||||
)
|
||||
|
||||
yield CompactStartEvent(
|
||||
current_context_tokens=old_tokens, threshold=threshold
|
||||
)
|
||||
|
||||
summary = await self.compact()
|
||||
|
||||
yield CompactEndEvent(
|
||||
old_context_tokens=old_tokens,
|
||||
new_context_tokens=self.stats.context_tokens,
|
||||
summary_length=len(summary),
|
||||
)
|
||||
|
||||
case MiddlewareAction.CONTINUE:
|
||||
pass
|
||||
|
||||
def _get_context(self) -> ConversationContext:
|
||||
return ConversationContext(
|
||||
messages=self.messages, stats=self.stats, config=self.config
|
||||
)
|
||||
|
||||
async def _conversation_loop(self, user_msg: str) -> AsyncGenerator[BaseEvent]:
|
||||
self.messages.append(LLMMessage(role=Role.user, content=user_msg))
|
||||
self.stats.steps += 1
|
||||
|
||||
try:
|
||||
should_break_loop = False
|
||||
while not should_break_loop:
|
||||
result = await self.middleware_pipeline.run_before_turn(
|
||||
self._get_context()
|
||||
)
|
||||
|
||||
async for event in self._handle_middleware_result(result):
|
||||
yield event
|
||||
|
||||
if result.action == MiddlewareAction.STOP:
|
||||
self._flush_new_messages()
|
||||
return
|
||||
|
||||
self.stats.steps += 1
|
||||
user_cancelled = False
|
||||
async for event in self._perform_llm_turn():
|
||||
if is_user_cancellation_event(event):
|
||||
user_cancelled = True
|
||||
yield event
|
||||
|
||||
last_message = self.messages[-1]
|
||||
should_break_loop = (
|
||||
last_message.role != Role.tool
|
||||
and self._last_chunk is not None
|
||||
and self._last_chunk.finish_reason is not None
|
||||
)
|
||||
|
||||
self._flush_new_messages()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
if user_cancelled:
|
||||
self._flush_new_messages()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
return
|
||||
|
||||
after_result = await self.middleware_pipeline.run_after_turn(
|
||||
self._get_context()
|
||||
)
|
||||
|
||||
async for event in self._handle_middleware_result(after_result):
|
||||
yield event
|
||||
|
||||
if after_result.action == MiddlewareAction.STOP:
|
||||
self._flush_new_messages()
|
||||
return
|
||||
|
||||
self._flush_new_messages()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
except Exception:
|
||||
self._flush_new_messages()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
raise
|
||||
|
||||
async def _perform_llm_turn(
|
||||
self,
|
||||
) -> AsyncGenerator[AssistantEvent | ToolCallEvent | ToolResultEvent]:
|
||||
if self.enable_streaming:
|
||||
async for event in self._stream_assistant_events():
|
||||
yield event
|
||||
else:
|
||||
assistant_event = await self._get_assistant_event()
|
||||
if assistant_event.content:
|
||||
yield assistant_event
|
||||
|
||||
last_message = self.messages[-1]
|
||||
last_chunk = self._last_chunk
|
||||
if last_chunk is None or last_chunk.usage is None:
|
||||
raise LLMResponseError("LLM response missing chunk or usage data")
|
||||
|
||||
parsed = self.format_handler.parse_message(last_message)
|
||||
resolved = self.format_handler.resolve_tool_calls(
|
||||
parsed, self.tool_manager, self.config
|
||||
)
|
||||
|
||||
if last_chunk.usage.completion_tokens > 0 and self.stats.last_turn_duration > 0:
|
||||
self.stats.tokens_per_second = (
|
||||
last_chunk.usage.completion_tokens / self.stats.last_turn_duration
|
||||
)
|
||||
|
||||
if not resolved.tool_calls and not resolved.failed_calls:
|
||||
return
|
||||
|
||||
async for event in self._handle_tool_calls(resolved):
|
||||
yield event
|
||||
|
||||
def _create_assistant_event(
|
||||
self, content: str, chunk: LLMChunk | None
|
||||
) -> AssistantEvent:
|
||||
return AssistantEvent(
|
||||
content=content,
|
||||
prompt_tokens=chunk.usage.prompt_tokens if chunk and chunk.usage else 0,
|
||||
completion_tokens=chunk.usage.completion_tokens
|
||||
if chunk and chunk.usage
|
||||
else 0,
|
||||
session_total_tokens=self.stats.session_total_llm_tokens,
|
||||
last_turn_duration=self.stats.last_turn_duration,
|
||||
tokens_per_second=self.stats.tokens_per_second,
|
||||
)
|
||||
|
||||
async def _stream_assistant_events(self) -> AsyncGenerator[AssistantEvent]:
|
||||
chunks: list[LLMChunk] = []
|
||||
content_buffer = ""
|
||||
chunks_with_content = 0
|
||||
BATCH_SIZE = 5
|
||||
|
||||
async for chunk in self._chat_streaming():
|
||||
chunks.append(chunk)
|
||||
|
||||
if chunk.message.tool_calls and chunk.finish_reason is None:
|
||||
if chunk.message.content:
|
||||
content_buffer += chunk.message.content
|
||||
chunks_with_content += 1
|
||||
|
||||
if content_buffer:
|
||||
yield self._create_assistant_event(content_buffer, chunk)
|
||||
content_buffer = ""
|
||||
chunks_with_content = 0
|
||||
continue
|
||||
|
||||
if chunk.message.content:
|
||||
content_buffer += chunk.message.content
|
||||
chunks_with_content += 1
|
||||
|
||||
if chunks_with_content >= BATCH_SIZE:
|
||||
yield self._create_assistant_event(content_buffer, chunk)
|
||||
content_buffer = ""
|
||||
chunks_with_content = 0
|
||||
|
||||
if content_buffer:
|
||||
last_chunk = chunks[-1] if chunks else None
|
||||
yield self._create_assistant_event(content_buffer, last_chunk)
|
||||
|
||||
full_content = ""
|
||||
full_tool_calls_map = OrderedDict[int, ToolCall]()
|
||||
for chunk in chunks:
|
||||
full_content += chunk.message.content or ""
|
||||
if not chunk.message.tool_calls:
|
||||
continue
|
||||
|
||||
for tc in chunk.message.tool_calls:
|
||||
if tc.index is None:
|
||||
raise LLMResponseError("Tool call chunk missing index")
|
||||
if tc.index not in full_tool_calls_map:
|
||||
full_tool_calls_map[tc.index] = tc
|
||||
else:
|
||||
new_args_str = (
|
||||
full_tool_calls_map[tc.index].function.arguments or ""
|
||||
) + (tc.function.arguments or "")
|
||||
full_tool_calls_map[tc.index].function.arguments = new_args_str
|
||||
|
||||
full_tool_calls = list(full_tool_calls_map.values()) or None
|
||||
last_message = LLMMessage(
|
||||
role=Role.assistant, content=full_content, tool_calls=full_tool_calls
|
||||
)
|
||||
self.messages.append(last_message)
|
||||
finish_reason = next(
|
||||
(c.finish_reason for c in chunks if c.finish_reason is not None), None
|
||||
)
|
||||
self._last_chunk = LLMChunk(
|
||||
message=last_message, usage=chunks[-1].usage, finish_reason=finish_reason
|
||||
)
|
||||
|
||||
async def _get_assistant_event(self) -> AssistantEvent:
|
||||
llm_result = await self._chat()
|
||||
if llm_result.usage is None:
|
||||
raise LLMResponseError(
|
||||
"Usage data missing in non-streaming completion response"
|
||||
)
|
||||
self._last_chunk = llm_result
|
||||
assistant_msg = llm_result.message
|
||||
self.messages.append(assistant_msg)
|
||||
|
||||
return AssistantEvent(
|
||||
content=assistant_msg.content or "",
|
||||
prompt_tokens=llm_result.usage.prompt_tokens,
|
||||
completion_tokens=llm_result.usage.completion_tokens,
|
||||
session_total_tokens=self.stats.session_total_llm_tokens,
|
||||
last_turn_duration=self.stats.last_turn_duration,
|
||||
tokens_per_second=self.stats.tokens_per_second,
|
||||
)
|
||||
|
||||
async def _handle_tool_calls( # noqa: PLR0915
|
||||
self, resolved: ResolvedMessage
|
||||
) -> AsyncGenerator[ToolCallEvent | ToolResultEvent]:
|
||||
for failed in resolved.failed_calls:
|
||||
error_msg = f"<{TOOL_ERROR_TAG}>{failed.tool_name}: {failed.error}</{TOOL_ERROR_TAG}>"
|
||||
|
||||
yield ToolResultEvent(
|
||||
tool_name=failed.tool_name,
|
||||
tool_class=None,
|
||||
error=error_msg,
|
||||
tool_call_id=failed.call_id,
|
||||
)
|
||||
|
||||
self.stats.tool_calls_failed += 1
|
||||
self.messages.append(
|
||||
self.format_handler.create_failed_tool_response_message(
|
||||
failed, error_msg
|
||||
)
|
||||
)
|
||||
|
||||
for tool_call in resolved.tool_calls:
|
||||
tool_call_id = tool_call.call_id
|
||||
|
||||
yield ToolCallEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
args=tool_call.validated_args,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
|
||||
try:
|
||||
tool_instance = self.tool_manager.get(tool_call.tool_name)
|
||||
except Exception as exc:
|
||||
error_msg = f"Error getting tool '{tool_call.tool_name}': {exc}"
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
error=error_msg,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, error_msg
|
||||
)
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
decision = await self._should_execute_tool(
|
||||
tool_instance, tool_call.args_dict, tool_call_id
|
||||
)
|
||||
|
||||
if decision.verdict == ToolExecutionResponse.SKIP:
|
||||
self.stats.tool_calls_rejected += 1
|
||||
skip_reason = decision.feedback or str(
|
||||
get_user_cancellation_message(
|
||||
CancellationReason.TOOL_SKIPPED, tool_call.tool_name
|
||||
)
|
||||
)
|
||||
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
skipped=True,
|
||||
skip_reason=skip_reason,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, skip_reason
|
||||
)
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
self.stats.tool_calls_agreed += 1
|
||||
|
||||
try:
|
||||
start_time = time.perf_counter()
|
||||
result_model = await tool_instance.invoke(**tool_call.args_dict)
|
||||
duration = time.perf_counter() - start_time
|
||||
|
||||
text = "\n".join(
|
||||
f"{k}: {v}" for k, v in result_model.model_dump().items()
|
||||
)
|
||||
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, text
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
result=result_model,
|
||||
duration=duration,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
|
||||
self.stats.tool_calls_succeeded += 1
|
||||
|
||||
except asyncio.CancelledError:
|
||||
cancel = str(
|
||||
get_user_cancellation_message(CancellationReason.TOOL_INTERRUPTED)
|
||||
)
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
error=cancel,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, cancel
|
||||
)
|
||||
)
|
||||
)
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
raise
|
||||
|
||||
except KeyboardInterrupt:
|
||||
cancel = str(
|
||||
get_user_cancellation_message(CancellationReason.TOOL_INTERRUPTED)
|
||||
)
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
error=cancel,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, cancel
|
||||
)
|
||||
)
|
||||
)
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
raise
|
||||
|
||||
except (ToolError, ToolPermissionError) as exc:
|
||||
error_msg = f"<{TOOL_ERROR_TAG}>{tool_instance.get_name()} failed: {exc}</{TOOL_ERROR_TAG}>"
|
||||
|
||||
yield ToolResultEvent(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_class=tool_call.tool_class,
|
||||
error=error_msg,
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
|
||||
if isinstance(exc, ToolPermissionError):
|
||||
self.stats.tool_calls_agreed -= 1
|
||||
self.stats.tool_calls_rejected += 1
|
||||
else:
|
||||
self.stats.tool_calls_failed += 1
|
||||
self.messages.append(
|
||||
LLMMessage.model_validate(
|
||||
self.format_handler.create_tool_response_message(
|
||||
tool_call, error_msg
|
||||
)
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
async def _chat(self, max_tokens: int | None = None) -> LLMChunk:
|
||||
active_model = self.config.get_active_model()
|
||||
provider = self.config.get_provider_for_model(active_model)
|
||||
|
||||
available_tools = self.format_handler.get_available_tools(
|
||||
self.tool_manager, self.config
|
||||
)
|
||||
tool_choice = self.format_handler.get_tool_choice()
|
||||
|
||||
try:
|
||||
start_time = time.perf_counter()
|
||||
|
||||
async with self.backend as backend:
|
||||
result = await backend.complete(
|
||||
model=active_model,
|
||||
messages=self.messages,
|
||||
temperature=active_model.temperature,
|
||||
tools=available_tools,
|
||||
tool_choice=tool_choice,
|
||||
extra_headers={
|
||||
"User-Agent": get_user_agent(),
|
||||
"x-affinity": self.session_id,
|
||||
},
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
end_time = time.perf_counter()
|
||||
if result.usage is None:
|
||||
raise LLMResponseError(
|
||||
"Usage data missing in non-streaming completion response"
|
||||
)
|
||||
|
||||
self.stats.last_turn_duration = end_time - start_time
|
||||
self.stats.last_turn_prompt_tokens = result.usage.prompt_tokens
|
||||
self.stats.last_turn_completion_tokens = result.usage.completion_tokens
|
||||
self.stats.session_prompt_tokens += result.usage.prompt_tokens
|
||||
self.stats.session_completion_tokens += result.usage.completion_tokens
|
||||
self.stats.context_tokens = (
|
||||
result.usage.prompt_tokens + result.usage.completion_tokens
|
||||
)
|
||||
|
||||
processed_message = self.format_handler.process_api_response_message(
|
||||
result.message
|
||||
)
|
||||
|
||||
return LLMChunk(
|
||||
message=processed_message,
|
||||
usage=result.usage,
|
||||
finish_reason=result.finish_reason,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"API error from {provider.name} (model: {active_model.name}): {e}"
|
||||
) from e
|
||||
|
||||
async def _chat_streaming(
|
||||
self, max_tokens: int | None = None
|
||||
) -> AsyncGenerator[LLMChunk]:
|
||||
active_model = self.config.get_active_model()
|
||||
provider = self.config.get_provider_for_model(active_model)
|
||||
|
||||
available_tools = self.format_handler.get_available_tools(
|
||||
self.tool_manager, self.config
|
||||
)
|
||||
tool_choice = self.format_handler.get_tool_choice()
|
||||
try:
|
||||
start_time = time.perf_counter()
|
||||
last_chunk = None
|
||||
async with self.backend as backend:
|
||||
async for chunk in backend.complete_streaming(
|
||||
model=active_model,
|
||||
messages=self.messages,
|
||||
temperature=active_model.temperature,
|
||||
tools=available_tools,
|
||||
tool_choice=tool_choice,
|
||||
extra_headers={
|
||||
"User-Agent": get_user_agent(),
|
||||
"x-affinity": self.session_id,
|
||||
},
|
||||
max_tokens=max_tokens,
|
||||
):
|
||||
last_chunk = chunk
|
||||
processed_message = (
|
||||
self.format_handler.process_api_response_message(chunk.message)
|
||||
)
|
||||
yield LLMChunk(
|
||||
message=processed_message,
|
||||
usage=chunk.usage,
|
||||
finish_reason=chunk.finish_reason,
|
||||
)
|
||||
|
||||
end_time = time.perf_counter()
|
||||
if last_chunk is None:
|
||||
raise LLMResponseError("Streamed completion returned no chunks")
|
||||
if last_chunk.usage is None:
|
||||
raise LLMResponseError(
|
||||
"Usage data missing in final chunk of streamed completion"
|
||||
)
|
||||
|
||||
self.stats.last_turn_duration = end_time - start_time
|
||||
self.stats.last_turn_prompt_tokens = last_chunk.usage.prompt_tokens
|
||||
self.stats.last_turn_completion_tokens = last_chunk.usage.completion_tokens
|
||||
self.stats.session_prompt_tokens += last_chunk.usage.prompt_tokens
|
||||
self.stats.session_completion_tokens += last_chunk.usage.completion_tokens
|
||||
self.stats.context_tokens = (
|
||||
last_chunk.usage.prompt_tokens + last_chunk.usage.completion_tokens
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"API error from {provider.name} (model: {active_model.name}): {e}"
|
||||
) from e
|
||||
|
||||
async def _should_execute_tool(
|
||||
self, tool: BaseTool, args: dict[str, Any], tool_call_id: str
|
||||
) -> ToolDecision:
|
||||
if self.auto_approve:
|
||||
return ToolDecision(verdict=ToolExecutionResponse.EXECUTE)
|
||||
|
||||
args_model, _ = tool._get_tool_args_results()
|
||||
validated_args = args_model.model_validate(args)
|
||||
|
||||
allowlist_denylist_result = tool.check_allowlist_denylist(validated_args)
|
||||
if allowlist_denylist_result == ToolPermission.ALWAYS:
|
||||
return ToolDecision(verdict=ToolExecutionResponse.EXECUTE)
|
||||
elif allowlist_denylist_result == ToolPermission.NEVER:
|
||||
denylist_patterns = tool.config.denylist
|
||||
denylist_str = ", ".join(repr(pattern) for pattern in denylist_patterns)
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.SKIP,
|
||||
feedback=f"Tool '{tool.get_name()}' blocked by denylist: [{denylist_str}]",
|
||||
)
|
||||
|
||||
tool_name = tool.get_name()
|
||||
perm = self.tool_manager.get_tool_config(tool_name).permission
|
||||
|
||||
if perm is ToolPermission.ALWAYS:
|
||||
return ToolDecision(verdict=ToolExecutionResponse.EXECUTE)
|
||||
if perm is ToolPermission.NEVER:
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.SKIP,
|
||||
feedback=f"Tool '{tool_name}' is permanently disabled",
|
||||
)
|
||||
|
||||
return await self._ask_approval(tool_name, args, tool_call_id)
|
||||
|
||||
async def _ask_approval(
|
||||
self, tool_name: str, args: dict[str, Any], tool_call_id: str
|
||||
) -> ToolDecision:
|
||||
if not self.approval_callback:
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.SKIP,
|
||||
feedback="Tool execution not permitted.",
|
||||
)
|
||||
if asyncio.iscoroutinefunction(self.approval_callback):
|
||||
response, feedback = await self.approval_callback(
|
||||
tool_name, args, tool_call_id
|
||||
)
|
||||
else:
|
||||
sync_callback = cast(SyncApprovalCallback, self.approval_callback)
|
||||
response, feedback = sync_callback(tool_name, args, tool_call_id)
|
||||
|
||||
match response:
|
||||
case ApprovalResponse.ALWAYS:
|
||||
self.auto_approve = True
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.EXECUTE, feedback=feedback
|
||||
)
|
||||
case ApprovalResponse.YES:
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.EXECUTE, feedback=feedback
|
||||
)
|
||||
case _:
|
||||
return ToolDecision(
|
||||
verdict=ToolExecutionResponse.SKIP, feedback=feedback
|
||||
)
|
||||
|
||||
def _clean_message_history(self) -> None:
|
||||
ACCEPTABLE_HISTORY_SIZE = 2
|
||||
if len(self.messages) < ACCEPTABLE_HISTORY_SIZE:
|
||||
return
|
||||
self._fill_missing_tool_responses()
|
||||
self._ensure_assistant_after_tools()
|
||||
|
||||
def _fill_missing_tool_responses(self) -> None:
|
||||
i = 1
|
||||
while i < len(self.messages): # noqa: PLR1702
|
||||
msg = self.messages[i]
|
||||
|
||||
if msg.role == "assistant" and msg.tool_calls:
|
||||
expected_responses = len(msg.tool_calls)
|
||||
|
||||
if expected_responses > 0:
|
||||
actual_responses = 0
|
||||
j = i + 1
|
||||
while j < len(self.messages) and self.messages[j].role == "tool":
|
||||
actual_responses += 1
|
||||
j += 1
|
||||
|
||||
if actual_responses < expected_responses:
|
||||
insertion_point = i + 1 + actual_responses
|
||||
|
||||
for call_idx in range(actual_responses, expected_responses):
|
||||
tool_call_data = msg.tool_calls[call_idx]
|
||||
|
||||
empty_response = LLMMessage(
|
||||
role=Role.tool,
|
||||
tool_call_id=tool_call_data.id or "",
|
||||
name=(tool_call_data.function.name or "")
|
||||
if tool_call_data.function
|
||||
else "",
|
||||
content=str(
|
||||
get_user_cancellation_message(
|
||||
CancellationReason.TOOL_NO_RESPONSE
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
self.messages.insert(insertion_point, empty_response)
|
||||
insertion_point += 1
|
||||
|
||||
i = i + 1 + expected_responses
|
||||
continue
|
||||
|
||||
i += 1
|
||||
|
||||
def _ensure_assistant_after_tools(self) -> None:
|
||||
MIN_MESSAGE_SIZE = 2
|
||||
if len(self.messages) < MIN_MESSAGE_SIZE:
|
||||
return
|
||||
|
||||
last_msg = self.messages[-1]
|
||||
if last_msg.role is Role.tool:
|
||||
empty_assistant_msg = LLMMessage(role=Role.assistant, content="Understood.")
|
||||
self.messages.append(empty_assistant_msg)
|
||||
|
||||
def _reset_session(self) -> None:
|
||||
self.session_id = str(uuid4())
|
||||
self.interaction_logger.reset_session(self.session_id)
|
||||
|
||||
def set_approval_callback(self, callback: ApprovalCallback) -> None:
|
||||
self.approval_callback = callback
|
||||
|
||||
async def clear_history(self) -> None:
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
self.messages = self.messages[:1]
|
||||
|
||||
self.stats = AgentStats()
|
||||
|
||||
try:
|
||||
active_model = self.config.get_active_model()
|
||||
self.stats.update_pricing(
|
||||
active_model.input_price, active_model.output_price
|
||||
)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
self.middleware_pipeline.reset()
|
||||
self.tool_manager.reset_all()
|
||||
self._reset_session()
|
||||
|
||||
async def compact(self) -> str:
|
||||
try:
|
||||
self._clean_message_history()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
last_user_message = None
|
||||
for msg in reversed(self.messages):
|
||||
if msg.role == Role.user:
|
||||
last_user_message = msg.content
|
||||
break
|
||||
|
||||
summary_request = UtilityPrompt.COMPACT.read()
|
||||
self.messages.append(LLMMessage(role=Role.user, content=summary_request))
|
||||
self.stats.steps += 1
|
||||
|
||||
summary_result = await self._chat()
|
||||
if summary_result.usage is None:
|
||||
raise LLMResponseError(
|
||||
"Usage data missing in compaction summary response"
|
||||
)
|
||||
summary_content = summary_result.message.content or ""
|
||||
|
||||
if last_user_message:
|
||||
summary_content += (
|
||||
f"\n\nLast request from user was: {last_user_message}"
|
||||
)
|
||||
|
||||
system_message = self.messages[0]
|
||||
summary_message = LLMMessage(role=Role.user, content=summary_content)
|
||||
self.messages = [system_message, summary_message]
|
||||
|
||||
active_model = self.config.get_active_model()
|
||||
|
||||
async with self.backend as backend:
|
||||
actual_context_tokens = await backend.count_tokens(
|
||||
model=active_model,
|
||||
messages=self.messages,
|
||||
tools=self.format_handler.get_available_tools(
|
||||
self.tool_manager, self.config
|
||||
),
|
||||
extra_headers={"User-Agent": get_user_agent()},
|
||||
)
|
||||
|
||||
self.stats.context_tokens = actual_context_tokens
|
||||
|
||||
self._reset_session()
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
self.middleware_pipeline.reset(reset_reason=ResetReason.COMPACT)
|
||||
|
||||
return summary_content or ""
|
||||
|
||||
except Exception:
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
raise
|
||||
|
||||
async def reload_with_initial_messages(
|
||||
self,
|
||||
config: VibeConfig | None = None,
|
||||
max_turns: int | None = None,
|
||||
max_price: float | None = None,
|
||||
) -> None:
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
|
||||
preserved_messages = self.messages[1:] if len(self.messages) > 1 else []
|
||||
old_system_prompt = self.messages[0].content if len(self.messages) > 0 else ""
|
||||
|
||||
if config is not None:
|
||||
self.config = config
|
||||
self.backend = self.backend_factory()
|
||||
|
||||
self.tool_manager = ToolManager(self.config)
|
||||
|
||||
new_system_prompt = get_universal_system_prompt(self.tool_manager, self.config)
|
||||
self.messages = [LLMMessage(role=Role.system, content=new_system_prompt)]
|
||||
did_system_prompt_change = old_system_prompt != new_system_prompt
|
||||
|
||||
if preserved_messages:
|
||||
self.messages.extend(preserved_messages)
|
||||
|
||||
if len(self.messages) == 1 or did_system_prompt_change:
|
||||
self.stats.reset_context_state()
|
||||
|
||||
try:
|
||||
active_model = self.config.get_active_model()
|
||||
self.stats.update_pricing(
|
||||
active_model.input_price, active_model.output_price
|
||||
)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
self._last_observed_message_index = 0
|
||||
|
||||
self._setup_middleware(max_turns, max_price)
|
||||
|
||||
if self.message_observer:
|
||||
for msg in self.messages:
|
||||
self.message_observer(msg)
|
||||
self._last_observed_message_index = len(self.messages)
|
||||
|
||||
self.tool_manager.reset_all()
|
||||
|
||||
await self.interaction_logger.save_interaction(
|
||||
self.messages, self.stats, self.config, self.tool_manager
|
||||
)
|
||||
0
vibe/core/autocompletion/__init__.py
Normal file
0
vibe/core/autocompletion/__init__.py
Normal file
247
vibe/core/autocompletion/completers.py
Normal file
247
vibe/core/autocompletion/completers.py
Normal file
|
|
@ -0,0 +1,247 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
from vibe.core.autocompletion.file_indexer import FileIndexer, IndexEntry
|
||||
from vibe.core.autocompletion.fuzzy import fuzzy_match
|
||||
|
||||
DEFAULT_MAX_ENTRIES_TO_PROCESS = 32000
|
||||
DEFAULT_TARGET_MATCHES = 100
|
||||
|
||||
|
||||
class Completer:
|
||||
def get_completions(self, text: str, cursor_pos: int) -> list[str]:
|
||||
return []
|
||||
|
||||
def get_completion_items(self, text: str, cursor_pos: int) -> list[tuple[str, str]]:
|
||||
return [
|
||||
(completion, "") for completion in self.get_completions(text, cursor_pos)
|
||||
]
|
||||
|
||||
def get_replacement_range(
|
||||
self, text: str, cursor_pos: int
|
||||
) -> tuple[int, int] | None:
|
||||
return None
|
||||
|
||||
|
||||
class CommandCompleter(Completer):
|
||||
def __init__(self, commands: list[tuple[str, str]]) -> None:
|
||||
aliases_with_descriptions: dict[str, str] = {}
|
||||
for alias, description in commands:
|
||||
aliases_with_descriptions[alias] = description
|
||||
|
||||
self._descriptions = aliases_with_descriptions
|
||||
self._aliases: list[str] = list(aliases_with_descriptions.keys())
|
||||
|
||||
def get_completions(self, text: str, cursor_pos: int) -> list[str]:
|
||||
if not text.startswith("/"):
|
||||
return []
|
||||
|
||||
word = text[1:cursor_pos].lower()
|
||||
search_str = "/" + word
|
||||
return [
|
||||
alias for alias in self._aliases if alias.lower().startswith(search_str)
|
||||
]
|
||||
|
||||
def get_completion_items(self, text: str, cursor_pos: int) -> list[tuple[str, str]]:
|
||||
completions = self.get_completions(text, cursor_pos)
|
||||
return [(alias, self._descriptions.get(alias, "")) for alias in completions]
|
||||
|
||||
def get_replacement_range(
|
||||
self, text: str, cursor_pos: int
|
||||
) -> tuple[int, int] | None:
|
||||
if text.startswith("/"):
|
||||
return (0, cursor_pos)
|
||||
return None
|
||||
|
||||
|
||||
class PathCompleter(Completer):
|
||||
def __init__(
|
||||
self,
|
||||
max_entries_to_process: int = DEFAULT_MAX_ENTRIES_TO_PROCESS,
|
||||
target_matches: int = DEFAULT_TARGET_MATCHES,
|
||||
) -> None:
|
||||
self._indexer = FileIndexer()
|
||||
self._max_entries_to_process = max_entries_to_process
|
||||
self._target_matches = target_matches
|
||||
|
||||
class _SearchContext(NamedTuple):
|
||||
suffix: str
|
||||
search_pattern: str
|
||||
path_prefix: str
|
||||
immediate_only: bool
|
||||
|
||||
def _extract_partial(self, before_cursor: str) -> str | None:
|
||||
if "@" not in before_cursor:
|
||||
return None
|
||||
|
||||
at_index = before_cursor.rfind("@")
|
||||
fragment = before_cursor[at_index + 1 :]
|
||||
|
||||
if " " in fragment:
|
||||
return None
|
||||
|
||||
return fragment
|
||||
|
||||
def _build_search_context(self, partial_path: str) -> _SearchContext:
|
||||
suffix = partial_path.split("/")[-1]
|
||||
|
||||
if not partial_path:
|
||||
# "@" => show top-level dir and files
|
||||
return self._SearchContext(
|
||||
search_pattern="", path_prefix="", suffix=suffix, immediate_only=True
|
||||
)
|
||||
|
||||
if partial_path.endswith("/"):
|
||||
# "@something/" => list immediate children
|
||||
return self._SearchContext(
|
||||
search_pattern="",
|
||||
path_prefix=partial_path,
|
||||
suffix=suffix,
|
||||
immediate_only=True,
|
||||
)
|
||||
|
||||
return self._SearchContext(
|
||||
# => run fuzzy search across the index
|
||||
search_pattern=partial_path,
|
||||
path_prefix="",
|
||||
suffix=suffix,
|
||||
immediate_only=False,
|
||||
)
|
||||
|
||||
def _matches_prefix(self, entry: IndexEntry, context: _SearchContext) -> bool:
|
||||
path_str = entry.rel
|
||||
|
||||
if context.path_prefix:
|
||||
prefix_without_slash = context.path_prefix.rstrip("/")
|
||||
prefix_with_slash = f"{prefix_without_slash}/"
|
||||
|
||||
if path_str == prefix_without_slash and entry.is_dir:
|
||||
# do not suggest the dir itself (e.g. "@src/" => don't suggest "@src/")
|
||||
return False
|
||||
|
||||
if path_str.startswith(prefix_with_slash):
|
||||
after_prefix = path_str[len(prefix_with_slash) :]
|
||||
else:
|
||||
idx = path_str.find(prefix_with_slash)
|
||||
if idx == -1 or (idx > 0 and path_str[idx - 1] != "/"):
|
||||
return False
|
||||
after_prefix = path_str[idx + len(prefix_with_slash) :]
|
||||
|
||||
# only suggest files/dirs that are immediate children of the prefix
|
||||
return bool(after_prefix) and "/" not in after_prefix
|
||||
|
||||
if context.immediate_only and "/" in path_str:
|
||||
# when user just typed "@", only show top-level entries
|
||||
return False
|
||||
|
||||
# entry matches the prefix: let the fuzzy matcher decide if it's a good match
|
||||
return True
|
||||
|
||||
def _is_visible(self, entry: IndexEntry, context: _SearchContext) -> bool:
|
||||
return not (entry.name.startswith(".") and not context.suffix.startswith("."))
|
||||
|
||||
def _format_label(self, entry: IndexEntry) -> str:
|
||||
suffix = "/" if entry.is_dir else ""
|
||||
return f"@{entry.rel}{suffix}"
|
||||
|
||||
def _score_matches(
|
||||
self, entries: list[IndexEntry], context: _SearchContext
|
||||
) -> list[tuple[str, float]]:
|
||||
scored_matches: list[tuple[str, float]] = []
|
||||
MAX_MATCHES = 50
|
||||
|
||||
for i, entry in enumerate(entries):
|
||||
if i >= self._max_entries_to_process:
|
||||
break
|
||||
|
||||
if not self._matches_prefix(entry, context):
|
||||
continue
|
||||
|
||||
if not self._is_visible(entry, context):
|
||||
continue
|
||||
|
||||
label = self._format_label(entry)
|
||||
|
||||
if not context.search_pattern:
|
||||
scored_matches.append((label, 0.0))
|
||||
if len(scored_matches) >= self._target_matches:
|
||||
break
|
||||
continue
|
||||
|
||||
match_result = fuzzy_match(
|
||||
context.search_pattern, entry.rel, entry.rel_lower
|
||||
)
|
||||
if match_result.matched:
|
||||
scored_matches.append((label, match_result.score))
|
||||
if (
|
||||
len(scored_matches) >= self._target_matches
|
||||
and match_result.score > MAX_MATCHES
|
||||
):
|
||||
break
|
||||
|
||||
scored_matches.sort(key=lambda x: (-x[1], x[0]))
|
||||
return scored_matches
|
||||
|
||||
def _collect_matches(self, text: str, cursor_pos: int) -> list[str]:
|
||||
before_cursor = text[:cursor_pos]
|
||||
partial_path = self._extract_partial(before_cursor)
|
||||
if partial_path is None:
|
||||
return []
|
||||
|
||||
context = self._build_search_context(partial_path)
|
||||
|
||||
try:
|
||||
# TODO (Vince): doing the assumption that "." is the root directory... Reliable?
|
||||
file_index = self._indexer.get_index(Path("."))
|
||||
except (OSError, RuntimeError):
|
||||
return []
|
||||
|
||||
scored_matches = self._score_matches(file_index, context)
|
||||
return [path for path, _ in scored_matches]
|
||||
|
||||
def get_completions(self, text: str, cursor_pos: int) -> list[str]:
|
||||
return self._collect_matches(text, cursor_pos)
|
||||
|
||||
def get_completion_items(self, text: str, cursor_pos: int) -> list[tuple[str, str]]:
|
||||
matches = self._collect_matches(text, cursor_pos)
|
||||
return [(completion, "") for completion in matches]
|
||||
|
||||
def get_replacement_range(
|
||||
self, text: str, cursor_pos: int
|
||||
) -> tuple[int, int] | None:
|
||||
before_cursor = text[:cursor_pos]
|
||||
if "@" in before_cursor:
|
||||
at_index = before_cursor.rfind("@")
|
||||
return (at_index, cursor_pos)
|
||||
return None
|
||||
|
||||
|
||||
class MultiCompleter(Completer):
|
||||
def __init__(self, completers: list[Completer]) -> None:
|
||||
self.completers = completers
|
||||
|
||||
def get_completions(self, text: str, cursor_pos: int) -> list[str]:
|
||||
all_completions = []
|
||||
for completer in self.completers:
|
||||
completions = completer.get_completions(text, cursor_pos)
|
||||
all_completions.extend(completions)
|
||||
|
||||
seen = set()
|
||||
unique = []
|
||||
for comp in all_completions:
|
||||
if comp not in seen:
|
||||
seen.add(comp)
|
||||
unique.append(comp)
|
||||
|
||||
return unique
|
||||
|
||||
def get_replacement_range(
|
||||
self, text: str, cursor_pos: int
|
||||
) -> tuple[int, int] | None:
|
||||
for completer in self.completers:
|
||||
range_result = completer.get_replacement_range(text, cursor_pos)
|
||||
if range_result is not None:
|
||||
return range_result
|
||||
return None
|
||||
10
vibe/core/autocompletion/file_indexer/__init__.py
Normal file
10
vibe/core/autocompletion/file_indexer/__init__.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.core.autocompletion.file_indexer.indexer import FileIndexer
|
||||
from vibe.core.autocompletion.file_indexer.store import (
|
||||
FileIndexStats,
|
||||
FileIndexStore,
|
||||
IndexEntry,
|
||||
)
|
||||
|
||||
__all__ = ["FileIndexStats", "FileIndexStore", "FileIndexer", "IndexEntry"]
|
||||
156
vibe/core/autocompletion/file_indexer/ignore_rules.py
Normal file
156
vibe/core/autocompletion/file_indexer/ignore_rules.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import fnmatch
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_IGNORE_PATTERNS: list[tuple[str, bool]] = [
|
||||
(".git/", True),
|
||||
("__pycache__/", True),
|
||||
("node_modules/", True),
|
||||
(".DS_Store", True),
|
||||
("*.pyc", True),
|
||||
("*.log", True),
|
||||
(".vscode/", True),
|
||||
(".idea/", True),
|
||||
("/build/", True),
|
||||
("dist/", True),
|
||||
("target/", True),
|
||||
(".next/", True),
|
||||
(".nuxt/", True),
|
||||
("coverage/", True),
|
||||
(".nyc_output/", True),
|
||||
("*.egg-info", True),
|
||||
(".pytest_cache/", True),
|
||||
(".tox/", True),
|
||||
("vendor/", True),
|
||||
("third_party/", True),
|
||||
("deps/", True),
|
||||
("*.min.js", True),
|
||||
("*.min.css", True),
|
||||
("*.bundle.js", True),
|
||||
("*.chunk.js", True),
|
||||
(".cache/", True),
|
||||
("tmp/", True),
|
||||
("temp/", True),
|
||||
("logs/", True),
|
||||
(".uv-cache/", True),
|
||||
(".ruff_cache/", True),
|
||||
(".venv/", True),
|
||||
("venv/", True),
|
||||
(".mypy_cache/", True),
|
||||
("htmlcov/", True),
|
||||
(".coverage", True),
|
||||
]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CompiledPattern:
|
||||
raw: str
|
||||
stripped: str
|
||||
is_exclude: bool
|
||||
dir_only: bool
|
||||
name_only: bool
|
||||
anchor_root: bool
|
||||
|
||||
|
||||
class IgnoreRules:
|
||||
def __init__(self, defaults: list[tuple[str, bool]] | None = None) -> None:
|
||||
self._defaults = defaults or DEFAULT_IGNORE_PATTERNS
|
||||
self._patterns: list[CompiledPattern] | None = None
|
||||
self._root: Path | None = None
|
||||
|
||||
def ensure_for_root(self, root: Path) -> None:
|
||||
resolved_root = root.resolve()
|
||||
if self._patterns is None or self._root != resolved_root:
|
||||
self._patterns = self._build_patterns(resolved_root)
|
||||
self._root = resolved_root
|
||||
|
||||
def should_ignore(self, rel_str: str, name: str, is_dir: bool) -> bool:
|
||||
if not self._patterns:
|
||||
return False
|
||||
|
||||
ignored = False
|
||||
for pattern in self._patterns:
|
||||
if self._matches(rel_str, name, is_dir, pattern):
|
||||
ignored = pattern.is_exclude
|
||||
return ignored
|
||||
|
||||
def reset(self) -> None:
|
||||
self._patterns = None
|
||||
self._root = None
|
||||
|
||||
def _build_patterns(self, root: Path) -> list[CompiledPattern]:
|
||||
patterns: list[CompiledPattern] = []
|
||||
for raw, is_exclude in self._defaults:
|
||||
anchor_root = raw.startswith("/")
|
||||
if anchor_root:
|
||||
raw = raw[1:]
|
||||
|
||||
stripped = raw.rstrip("/")
|
||||
patterns.append(
|
||||
CompiledPattern(
|
||||
raw=raw,
|
||||
stripped=stripped,
|
||||
is_exclude=is_exclude,
|
||||
dir_only=raw.endswith("/"),
|
||||
name_only="/" not in stripped,
|
||||
anchor_root=anchor_root,
|
||||
)
|
||||
)
|
||||
|
||||
gitignore_path = root / ".gitignore"
|
||||
if gitignore_path.exists():
|
||||
try:
|
||||
text = gitignore_path.read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
return patterns
|
||||
|
||||
for line in text.splitlines():
|
||||
raw = line.strip()
|
||||
if not raw or raw.startswith("#"):
|
||||
continue
|
||||
|
||||
if "#" in raw:
|
||||
raw = raw.split("#", 1)[0].rstrip()
|
||||
if not raw:
|
||||
continue
|
||||
|
||||
is_exclude = not raw.startswith("!")
|
||||
if not is_exclude:
|
||||
raw = raw[1:].lstrip()
|
||||
if not raw:
|
||||
continue
|
||||
|
||||
anchor_root = raw.startswith("/")
|
||||
if anchor_root:
|
||||
raw = raw[1:]
|
||||
|
||||
stripped = raw.rstrip("/")
|
||||
patterns.append(
|
||||
CompiledPattern(
|
||||
raw=raw,
|
||||
stripped=stripped,
|
||||
is_exclude=is_exclude,
|
||||
dir_only=raw.endswith("/"),
|
||||
name_only="/" not in stripped,
|
||||
anchor_root=anchor_root,
|
||||
)
|
||||
)
|
||||
|
||||
return patterns
|
||||
|
||||
def _matches(
|
||||
self, rel_str: str, name: str, is_dir: bool, pattern: CompiledPattern
|
||||
) -> bool:
|
||||
if pattern.name_only:
|
||||
if pattern.anchor_root and "/" in rel_str:
|
||||
return False
|
||||
target = name
|
||||
else:
|
||||
target = rel_str
|
||||
|
||||
if not fnmatch.fnmatch(target, pattern.stripped):
|
||||
return False
|
||||
|
||||
return not pattern.dir_only or is_dir
|
||||
176
vibe/core/autocompletion/file_indexer/indexer.py
Normal file
176
vibe/core/autocompletion/file_indexer/indexer.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from threading import Event, RLock
|
||||
|
||||
from vibe.core.autocompletion.file_indexer.ignore_rules import IgnoreRules
|
||||
from vibe.core.autocompletion.file_indexer.store import (
|
||||
FileIndexStats,
|
||||
FileIndexStore,
|
||||
IndexEntry,
|
||||
)
|
||||
from vibe.core.autocompletion.file_indexer.watcher import Change, WatchController
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _RebuildTask:
|
||||
cancel_event: Event
|
||||
done_event: Event
|
||||
|
||||
|
||||
class FileIndexer:
|
||||
def __init__(self, mass_change_threshold: int = 200) -> None:
|
||||
self._lock = RLock() # guards _store snapshot access and watcher callbacks.
|
||||
self._stats = FileIndexStats()
|
||||
self._ignore_rules = IgnoreRules()
|
||||
self._store = FileIndexStore(
|
||||
self._ignore_rules, self._stats, mass_change_threshold=mass_change_threshold
|
||||
)
|
||||
self._watcher = WatchController(self._handle_watch_changes)
|
||||
self._rebuild_executor = ThreadPoolExecutor(
|
||||
max_workers=1, thread_name_prefix="file-indexer"
|
||||
)
|
||||
self._active_rebuilds: dict[Path, _RebuildTask] = {}
|
||||
self._rebuild_lock = (
|
||||
RLock()
|
||||
) # coordinates updates to _active_rebuilds and _target_root.
|
||||
self._target_root: Path | None = None
|
||||
self._shutdown = False
|
||||
|
||||
@property
|
||||
def stats(self) -> FileIndexStats:
|
||||
return self._stats
|
||||
|
||||
def get_index(self, root: Path) -> list[IndexEntry]:
|
||||
resolved_root = root.resolve()
|
||||
|
||||
with self._lock: # read current root without blocking rebuild bookkeeping
|
||||
root_changed = (
|
||||
self._store.root is not None and self._store.root != resolved_root
|
||||
)
|
||||
|
||||
if root_changed:
|
||||
self._watcher.stop()
|
||||
with self._rebuild_lock: # cancel rebuilds targeting other roots
|
||||
self._target_root = resolved_root
|
||||
for other_root, task in self._active_rebuilds.items():
|
||||
if other_root != resolved_root:
|
||||
task.cancel_event.set()
|
||||
task.done_event.set()
|
||||
self._active_rebuilds.pop(other_root, None)
|
||||
|
||||
with self._lock:
|
||||
needs_rebuild = self._store.root != resolved_root
|
||||
|
||||
if needs_rebuild:
|
||||
with self._rebuild_lock:
|
||||
self._target_root = resolved_root
|
||||
self._start_background_rebuild(resolved_root)
|
||||
self._wait_for_rebuild(resolved_root)
|
||||
|
||||
self._watcher.start(resolved_root)
|
||||
|
||||
with self._lock: # ensure root reference is fresh before snapshotting
|
||||
return self._store.snapshot()
|
||||
|
||||
def refresh(self) -> None:
|
||||
self._watcher.stop()
|
||||
with self._rebuild_lock:
|
||||
for task in self._active_rebuilds.values():
|
||||
task.cancel_event.set()
|
||||
task.done_event.set()
|
||||
self._active_rebuilds.clear()
|
||||
self._target_root = None
|
||||
with self._lock:
|
||||
self._store.clear()
|
||||
self._ignore_rules.reset()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self._shutdown:
|
||||
return
|
||||
self._shutdown = True
|
||||
self.refresh()
|
||||
self._rebuild_executor.shutdown(wait=True)
|
||||
|
||||
def __del__(self) -> None:
|
||||
if not self._shutdown:
|
||||
try:
|
||||
self.shutdown()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _start_background_rebuild(self, root: Path) -> None:
|
||||
with self._rebuild_lock: # one rebuild per root
|
||||
if root in self._active_rebuilds:
|
||||
return
|
||||
|
||||
cancel_event = Event()
|
||||
done_event = Event()
|
||||
self._active_rebuilds[root] = _RebuildTask(
|
||||
cancel_event=cancel_event, done_event=done_event
|
||||
)
|
||||
|
||||
try:
|
||||
self._rebuild_executor.submit(
|
||||
self._rebuild_worker, root, self._active_rebuilds[root]
|
||||
)
|
||||
except RuntimeError:
|
||||
with self._rebuild_lock:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
done_event.set()
|
||||
|
||||
def _rebuild_worker(self, root: Path, task: _RebuildTask) -> None:
|
||||
try:
|
||||
if task.cancel_event.is_set(): # cancelled before work began
|
||||
with self._rebuild_lock:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
return
|
||||
|
||||
with self._rebuild_lock: # bail if another root took ownership
|
||||
if self._target_root != root:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
return
|
||||
|
||||
with self._lock: # exclusive access while rebuilding the store
|
||||
if task.cancel_event.is_set():
|
||||
with self._rebuild_lock:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
return
|
||||
|
||||
self._store.rebuild(
|
||||
root, should_cancel=lambda: task.cancel_event.is_set()
|
||||
)
|
||||
|
||||
with self._rebuild_lock:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
except Exception:
|
||||
with self._rebuild_lock:
|
||||
self._active_rebuilds.pop(root, None)
|
||||
finally:
|
||||
task.done_event.set()
|
||||
|
||||
def _wait_for_rebuild(self, root: Path) -> None:
|
||||
with self._rebuild_lock:
|
||||
task = self._active_rebuilds.get(root)
|
||||
if task:
|
||||
task.done_event.wait()
|
||||
|
||||
def _handle_watch_changes(
|
||||
self, root: Path, raw_changes: Iterable[tuple[Change, str]]
|
||||
) -> None:
|
||||
normalized: list[tuple[Change, Path]] = []
|
||||
for change, path_str in raw_changes:
|
||||
if change not in {Change.added, Change.deleted, Change.modified}:
|
||||
continue
|
||||
normalized.append((change, Path(path_str).resolve()))
|
||||
|
||||
if not normalized:
|
||||
return
|
||||
|
||||
with self._lock: # make watcher ignore stale roots
|
||||
if self._store.root != root:
|
||||
return
|
||||
self._store.apply_changes(normalized)
|
||||
169
vibe/core/autocompletion/file_indexer/store.py
Normal file
169
vibe/core/autocompletion/file_indexer/store.py
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from vibe.core.autocompletion.file_indexer.ignore_rules import IgnoreRules
|
||||
from vibe.core.autocompletion.file_indexer.watcher import Change
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FileIndexStats:
|
||||
rebuilds: int = 0
|
||||
incremental_updates: int = 0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexEntry:
|
||||
rel: str
|
||||
rel_lower: str
|
||||
name: str
|
||||
path: Path
|
||||
is_dir: bool
|
||||
|
||||
|
||||
class FileIndexStore:
|
||||
def __init__(
|
||||
self,
|
||||
ignore_rules: IgnoreRules,
|
||||
stats: FileIndexStats,
|
||||
mass_change_threshold: int = 200,
|
||||
) -> None:
|
||||
self._ignore_rules = ignore_rules
|
||||
self._stats = stats
|
||||
self._mass_change_threshold = mass_change_threshold
|
||||
self._entries_by_rel: dict[str, IndexEntry] = {}
|
||||
self._ordered_entries: list[IndexEntry] | None = None
|
||||
self._root: Path | None = None
|
||||
|
||||
@property
|
||||
def root(self) -> Path | None:
|
||||
return self._root
|
||||
|
||||
def clear(self) -> None:
|
||||
self._entries_by_rel.clear()
|
||||
self._ordered_entries = None
|
||||
self._root = None
|
||||
|
||||
def rebuild(
|
||||
self, root: Path, should_cancel: Callable[[], bool] | None = None
|
||||
) -> None:
|
||||
resolved_root = root.resolve()
|
||||
self._ignore_rules.ensure_for_root(resolved_root)
|
||||
entries = self._walk_directory(resolved_root, cancel_check=should_cancel)
|
||||
self._entries_by_rel = {entry.rel: entry for entry in entries}
|
||||
self._ordered_entries = entries
|
||||
self._root = resolved_root
|
||||
self._stats.rebuilds += 1
|
||||
|
||||
def snapshot(self) -> list[IndexEntry]:
|
||||
if not self._entries_by_rel:
|
||||
return []
|
||||
|
||||
if self._ordered_entries is None:
|
||||
self._ordered_entries = sorted(
|
||||
self._entries_by_rel.values(), key=lambda entry: entry.rel
|
||||
)
|
||||
|
||||
return list(self._ordered_entries)
|
||||
|
||||
def apply_changes(self, changes: list[tuple[Change, Path]]) -> None:
|
||||
if self._root is None:
|
||||
return
|
||||
|
||||
if len(changes) > self._mass_change_threshold:
|
||||
self.rebuild(self._root)
|
||||
return
|
||||
|
||||
modified = False
|
||||
for change, path in changes:
|
||||
try:
|
||||
rel_str = path.relative_to(self._root).as_posix()
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if not rel_str:
|
||||
continue
|
||||
|
||||
if change is Change.deleted:
|
||||
if self._remove_entry(rel_str):
|
||||
modified = True
|
||||
continue
|
||||
|
||||
if not path.exists():
|
||||
continue
|
||||
|
||||
if path.is_dir():
|
||||
dir_entry = self._create_entry(rel_str, path.name, path, True)
|
||||
if dir_entry:
|
||||
self._entries_by_rel[rel_str] = dir_entry
|
||||
modified = True
|
||||
for entry in self._walk_directory(path, rel_str):
|
||||
self._entries_by_rel[entry.rel] = entry
|
||||
modified = True
|
||||
else:
|
||||
file_entry = self._create_entry(rel_str, path.name, path, False)
|
||||
if file_entry:
|
||||
self._entries_by_rel[file_entry.rel] = file_entry
|
||||
modified = True
|
||||
|
||||
if modified:
|
||||
self._ordered_entries = None
|
||||
self._stats.incremental_updates += 1
|
||||
|
||||
def _create_entry(
|
||||
self, rel_str: str, name: str, path: Path, is_dir: bool
|
||||
) -> IndexEntry | None:
|
||||
if self._ignore_rules.should_ignore(rel_str, name, is_dir):
|
||||
return None
|
||||
return IndexEntry(
|
||||
rel=rel_str, rel_lower=rel_str.lower(), name=name, path=path, is_dir=is_dir
|
||||
)
|
||||
|
||||
def _walk_directory(
|
||||
self,
|
||||
directory: Path,
|
||||
rel_prefix: str = "",
|
||||
cancel_check: Callable[[], bool] | None = None,
|
||||
) -> list[IndexEntry]:
|
||||
results: list[IndexEntry] = []
|
||||
try:
|
||||
with os.scandir(directory) as iterator:
|
||||
for entry in iterator:
|
||||
if cancel_check and cancel_check():
|
||||
break
|
||||
|
||||
is_dir = entry.is_dir(follow_symlinks=False)
|
||||
name = entry.name
|
||||
rel_str = f"{rel_prefix}/{name}" if rel_prefix else name
|
||||
path = Path(entry.path)
|
||||
|
||||
index_entry = self._create_entry(rel_str, name, path, is_dir)
|
||||
if not index_entry:
|
||||
continue
|
||||
|
||||
results.append(index_entry)
|
||||
|
||||
if is_dir:
|
||||
results.extend(
|
||||
self._walk_directory(path, rel_str, cancel_check)
|
||||
)
|
||||
except (PermissionError, OSError):
|
||||
pass
|
||||
|
||||
return results
|
||||
|
||||
def _remove_entry(self, rel_str: str) -> bool:
|
||||
entry = self._entries_by_rel.pop(rel_str, None)
|
||||
if not entry:
|
||||
return False
|
||||
|
||||
if entry.is_dir:
|
||||
prefix = f"{rel_str}/"
|
||||
to_remove = [key for key in self._entries_by_rel if key.startswith(prefix)]
|
||||
for key in to_remove:
|
||||
self._entries_by_rel.pop(key, None)
|
||||
|
||||
return True
|
||||
71
vibe/core/autocompletion/file_indexer/watcher.py
Normal file
71
vibe/core/autocompletion/file_indexer/watcher.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from pathlib import Path
|
||||
from threading import Event, Thread
|
||||
|
||||
from watchfiles import Change, watch
|
||||
|
||||
|
||||
class WatchController:
|
||||
def __init__(
|
||||
self, on_changes: Callable[[Path, Iterable[tuple[Change, str]]], None]
|
||||
) -> None:
|
||||
self._on_changes = on_changes
|
||||
self._thread: Thread | None = None
|
||||
self._stop_event: Event | None = None
|
||||
self._ready_event: Event | None = None
|
||||
self._root: Path | None = None
|
||||
|
||||
def start(self, root: Path) -> None:
|
||||
resolved_root = root.resolve()
|
||||
if self._thread and self._thread.is_alive() and self._root == resolved_root:
|
||||
return
|
||||
|
||||
self.stop()
|
||||
|
||||
stop_event = Event()
|
||||
ready_event = Event()
|
||||
thread = Thread(
|
||||
target=self._watch_loop,
|
||||
args=(resolved_root, stop_event, ready_event),
|
||||
name="file-indexer-watch",
|
||||
daemon=True,
|
||||
)
|
||||
|
||||
self._thread = thread
|
||||
self._stop_event = stop_event
|
||||
self._ready_event = ready_event
|
||||
self._root = resolved_root
|
||||
|
||||
thread.start()
|
||||
ready_event.wait(timeout=0.5)
|
||||
|
||||
def stop(self) -> None:
|
||||
thread = self._thread
|
||||
if self._stop_event:
|
||||
self._stop_event.set()
|
||||
self._thread = None
|
||||
self._stop_event = None
|
||||
self._ready_event = None
|
||||
self._root = None
|
||||
|
||||
if thread and thread.is_alive():
|
||||
thread.join(timeout=1)
|
||||
|
||||
def _watch_loop(self, root: Path, stop_event: Event, ready_event: Event) -> None:
|
||||
try:
|
||||
watcher = watch(
|
||||
str(root), stop_event=stop_event, step=200, yield_on_timeout=True
|
||||
)
|
||||
ready_event.set()
|
||||
for changes in watcher:
|
||||
if not ready_event.is_set():
|
||||
ready_event.set()
|
||||
if stop_event.is_set():
|
||||
break
|
||||
if not changes:
|
||||
continue
|
||||
self._on_changes(root, changes)
|
||||
except Exception:
|
||||
ready_event.set()
|
||||
189
vibe/core/autocompletion/fuzzy.py
Normal file
189
vibe/core/autocompletion/fuzzy.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
PREFIX_MULTIPLIER = 2.0
|
||||
WORD_BOUNDARY_MULTIPLIER = 1.8
|
||||
CONSECUTIVE_MULTIPLIER = 1.3
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MatchResult:
|
||||
matched: bool
|
||||
score: float
|
||||
matched_indices: tuple[int, ...]
|
||||
|
||||
|
||||
def fuzzy_match(pattern: str, text: str, text_lower: str | None = None) -> MatchResult:
|
||||
if not pattern:
|
||||
return MatchResult(matched=True, score=0.0, matched_indices=())
|
||||
|
||||
if text_lower is None:
|
||||
text_lower = text.lower()
|
||||
return _find_best_match(pattern, pattern.lower(), text_lower, text)
|
||||
|
||||
|
||||
def _find_best_match(
|
||||
pattern_original: str, pattern_lower: str, text_lower: str, text_original: str
|
||||
) -> MatchResult:
|
||||
if len(pattern_lower) > len(text_lower):
|
||||
return MatchResult(matched=False, score=0.0, matched_indices=())
|
||||
|
||||
if text_lower.startswith(pattern_lower):
|
||||
indices = tuple(range(len(pattern_lower)))
|
||||
score = _calculate_score(
|
||||
pattern_original, pattern_lower, text_lower, indices, text_original
|
||||
)
|
||||
return MatchResult(
|
||||
matched=True, score=score * PREFIX_MULTIPLIER, matched_indices=indices
|
||||
)
|
||||
|
||||
best_score = -1.0
|
||||
best_indices: tuple[int, ...] = ()
|
||||
|
||||
for matcher in (
|
||||
_try_word_boundary_match,
|
||||
_try_consecutive_match,
|
||||
_try_subsequence_match,
|
||||
):
|
||||
match = matcher(pattern_original, pattern_lower, text_lower, text_original)
|
||||
if match.matched and match.score > best_score:
|
||||
best_score = match.score
|
||||
best_indices = match.matched_indices
|
||||
|
||||
if best_score >= 0:
|
||||
return MatchResult(matched=True, score=best_score, matched_indices=best_indices)
|
||||
|
||||
return MatchResult(matched=False, score=0.0, matched_indices=())
|
||||
|
||||
|
||||
def _try_word_boundary_match(
|
||||
pattern_original: str, pattern: str, text_lower: str, text_original: str
|
||||
) -> MatchResult:
|
||||
indices: list[int] = []
|
||||
pattern_idx = 0
|
||||
|
||||
for i, char in enumerate(text_lower):
|
||||
if pattern_idx >= len(pattern):
|
||||
break
|
||||
|
||||
is_boundary = (
|
||||
i == 0
|
||||
or text_lower[i - 1] in "/-_."
|
||||
or (text_original[i].isupper() and not text_original[i - 1].isupper())
|
||||
)
|
||||
|
||||
if char == pattern[pattern_idx]:
|
||||
if is_boundary or (indices and i == indices[-1] + 1) or not indices:
|
||||
indices.append(i)
|
||||
pattern_idx += 1
|
||||
|
||||
if pattern_idx == len(pattern):
|
||||
score = _calculate_score(
|
||||
pattern_original, pattern, text_lower, tuple(indices), text_original
|
||||
)
|
||||
return MatchResult(
|
||||
matched=True,
|
||||
score=score * WORD_BOUNDARY_MULTIPLIER,
|
||||
matched_indices=tuple(indices),
|
||||
)
|
||||
|
||||
return MatchResult(matched=False, score=0.0, matched_indices=())
|
||||
|
||||
|
||||
def _try_consecutive_match(
|
||||
pattern_original: str, pattern: str, text_lower: str, text_original: str
|
||||
) -> MatchResult:
|
||||
indices: list[int] = []
|
||||
pattern_idx = 0
|
||||
|
||||
for i, char in enumerate(text_lower):
|
||||
if pattern_idx >= len(pattern):
|
||||
break
|
||||
|
||||
if char == pattern[pattern_idx]:
|
||||
indices.append(i)
|
||||
pattern_idx += 1
|
||||
elif indices:
|
||||
indices.clear()
|
||||
pattern_idx = 0
|
||||
|
||||
if pattern_idx == len(pattern):
|
||||
score = _calculate_score(
|
||||
pattern_original, pattern, text_lower, tuple(indices), text_original
|
||||
)
|
||||
return MatchResult(
|
||||
matched=True,
|
||||
score=score * CONSECUTIVE_MULTIPLIER,
|
||||
matched_indices=tuple(indices),
|
||||
)
|
||||
|
||||
return MatchResult(matched=False, score=0.0, matched_indices=())
|
||||
|
||||
|
||||
def _try_subsequence_match(
|
||||
pattern_original: str, pattern: str, text_lower: str, text_original: str
|
||||
) -> MatchResult:
|
||||
indices: list[int] = []
|
||||
pattern_idx = 0
|
||||
|
||||
for i, char in enumerate(text_lower):
|
||||
if pattern_idx >= len(pattern):
|
||||
break
|
||||
if char == pattern[pattern_idx]:
|
||||
indices.append(i)
|
||||
pattern_idx += 1
|
||||
|
||||
if pattern_idx == len(pattern):
|
||||
score = _calculate_score(
|
||||
pattern_original, pattern, text_lower, tuple(indices), text_original
|
||||
)
|
||||
return MatchResult(matched=True, score=score, matched_indices=tuple(indices))
|
||||
|
||||
return MatchResult(matched=False, score=0.0, matched_indices=())
|
||||
|
||||
|
||||
def _calculate_score(
|
||||
pattern_original: str,
|
||||
pattern: str,
|
||||
text_lower: str,
|
||||
indices: tuple[int, ...],
|
||||
text_original: str,
|
||||
) -> float:
|
||||
if not indices:
|
||||
return 0.0
|
||||
|
||||
base_score = 100.0
|
||||
if indices[0] == 0:
|
||||
base_score += 50.0
|
||||
else:
|
||||
base_score -= indices[0] * 2
|
||||
|
||||
consecutive_bonus = sum(
|
||||
10.0 for i in range(len(indices) - 1) if indices[i + 1] == indices[i] + 1
|
||||
)
|
||||
|
||||
boundary_bonus = 0.0
|
||||
for idx in indices:
|
||||
if idx == 0 or text_lower[idx - 1] in "/-_.":
|
||||
boundary_bonus += 5.0
|
||||
elif text_original[idx].isupper() and (
|
||||
idx == 0 or not text_original[idx - 1].isupper()
|
||||
):
|
||||
boundary_bonus += 3.0
|
||||
|
||||
case_bonus = sum(
|
||||
2.0
|
||||
for i, text_idx in enumerate(indices)
|
||||
if i < len(pattern_original)
|
||||
and text_idx < len(text_original)
|
||||
and pattern_original[i] == text_original[text_idx]
|
||||
)
|
||||
|
||||
gap_penalty = sum(
|
||||
(indices[i + 1] - indices[i] - 1) * 1.5 for i in range(len(indices) - 1)
|
||||
)
|
||||
|
||||
return max(
|
||||
0.0, base_score + consecutive_bonus + boundary_bonus + case_bonus - gap_penalty
|
||||
)
|
||||
108
vibe/core/autocompletion/path_prompt.py
Normal file
108
vibe/core/autocompletion/path_prompt.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PathResource:
|
||||
path: Path
|
||||
alias: str
|
||||
kind: Literal["file", "directory"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PathPromptPayload:
|
||||
display_text: str
|
||||
prompt_text: str
|
||||
resources: list[PathResource]
|
||||
|
||||
|
||||
def build_path_prompt_payload(
|
||||
message: str, *, base_dir: Path | None = None
|
||||
) -> PathPromptPayload:
|
||||
if not message:
|
||||
return PathPromptPayload(message, message, [])
|
||||
|
||||
resolved_base = (base_dir or Path.cwd()).resolve()
|
||||
prompt_parts: list[str] = []
|
||||
resources: list[PathResource] = []
|
||||
pos = 0
|
||||
|
||||
while pos < len(message):
|
||||
if _is_path_anchor(message, pos):
|
||||
candidate, new_pos = _extract_candidate(message, pos + 1)
|
||||
if candidate and (resource := _to_resource(candidate, resolved_base)):
|
||||
resources.append(resource)
|
||||
prompt_parts.append(candidate)
|
||||
pos = new_pos
|
||||
continue
|
||||
|
||||
prompt_parts.append(message[pos])
|
||||
pos += 1
|
||||
|
||||
prompt_text = "".join(prompt_parts)
|
||||
unique_resources = _dedupe_resources(resources)
|
||||
return PathPromptPayload(message, prompt_text, unique_resources)
|
||||
|
||||
|
||||
def _is_path_anchor(message: str, pos: int) -> bool:
|
||||
if message[pos] != "@":
|
||||
return False
|
||||
if pos == 0:
|
||||
return True
|
||||
return not (message[pos - 1].isalnum() or message[pos - 1] == "_")
|
||||
|
||||
|
||||
def _extract_candidate(message: str, start: int) -> tuple[str | None, int]:
|
||||
if start >= len(message):
|
||||
return None, start
|
||||
|
||||
quote = message[start]
|
||||
if quote in {"'", '"'}:
|
||||
end_quote = message.find(quote, start + 1)
|
||||
if end_quote == -1:
|
||||
return None, start
|
||||
return message[start + 1 : end_quote], end_quote + 1
|
||||
|
||||
end = start
|
||||
while end < len(message) and _is_path_char(message[end]):
|
||||
end += 1
|
||||
|
||||
if end == start:
|
||||
return None, start
|
||||
|
||||
return message[start:end], end
|
||||
|
||||
|
||||
def _is_path_char(char: str) -> bool:
|
||||
return char.isalnum() or char in "._/\\-()[]{}"
|
||||
|
||||
|
||||
def _to_resource(candidate: str, base_dir: Path) -> PathResource | None:
|
||||
if not candidate:
|
||||
return None
|
||||
|
||||
candidate_path = Path(candidate)
|
||||
resolved = (
|
||||
candidate_path if candidate_path.is_absolute() else base_dir / candidate_path
|
||||
)
|
||||
resolved = resolved.resolve()
|
||||
|
||||
if not resolved.exists():
|
||||
return None
|
||||
|
||||
kind = "directory" if resolved.is_dir() else "file"
|
||||
return PathResource(path=resolved, alias=candidate, kind=kind)
|
||||
|
||||
|
||||
def _dedupe_resources(resources: list[PathResource]) -> list[PathResource]:
|
||||
seen: set[Path] = set()
|
||||
unique: list[PathResource] = []
|
||||
for resource in resources:
|
||||
if resource.path in seen:
|
||||
continue
|
||||
seen.add(resource.path)
|
||||
unique.append(resource)
|
||||
return unique
|
||||
149
vibe/core/autocompletion/path_prompt_adapter.py
Normal file
149
vibe/core/autocompletion/path_prompt_adapter.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
import mimetypes
|
||||
from pathlib import Path
|
||||
|
||||
from vibe.core.autocompletion.path_prompt import (
|
||||
PathPromptPayload,
|
||||
PathResource,
|
||||
build_path_prompt_payload,
|
||||
)
|
||||
|
||||
DEFAULT_MAX_EMBED_BYTES = 256 * 1024
|
||||
|
||||
ResourceBlock = dict[str, str | None]
|
||||
|
||||
|
||||
def render_path_prompt(
|
||||
message: str,
|
||||
*,
|
||||
base_dir: Path,
|
||||
max_embed_bytes: int | None = DEFAULT_MAX_EMBED_BYTES,
|
||||
) -> str:
|
||||
payload = build_path_prompt_payload(message, base_dir=base_dir)
|
||||
blocks = _path_prompt_to_content_blocks(payload, max_embed_bytes=max_embed_bytes)
|
||||
return _content_blocks_to_prompt_text(blocks)
|
||||
|
||||
|
||||
def _path_prompt_to_content_blocks(
|
||||
payload: PathPromptPayload, *, max_embed_bytes: int | None = DEFAULT_MAX_EMBED_BYTES
|
||||
) -> list[ResourceBlock]:
|
||||
blocks: list[ResourceBlock] = [{"type": "text", "text": payload.prompt_text}]
|
||||
|
||||
for resource in payload.resources:
|
||||
match resource.kind:
|
||||
case "file":
|
||||
embedded = _try_embed_text_resource(resource, max_embed_bytes)
|
||||
if embedded:
|
||||
blocks.append(embedded)
|
||||
else:
|
||||
blocks.append({
|
||||
"type": "resource_link",
|
||||
"uri": resource.path.as_uri(),
|
||||
"name": resource.alias,
|
||||
})
|
||||
case "directory":
|
||||
blocks.append({
|
||||
"type": "resource_link",
|
||||
"uri": resource.path.as_uri(),
|
||||
"name": resource.alias,
|
||||
})
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
def _try_embed_text_resource(
|
||||
resource: PathResource, max_embed_bytes: int | None
|
||||
) -> ResourceBlock | None:
|
||||
try:
|
||||
data = resource.path.read_bytes()
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
if max_embed_bytes is not None and len(data) > max_embed_bytes:
|
||||
return None
|
||||
|
||||
if not _is_probably_text(resource, data):
|
||||
return None
|
||||
|
||||
try:
|
||||
text = data.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
|
||||
return {"type": "resource", "uri": resource.path.as_uri(), "text": text}
|
||||
|
||||
|
||||
def _content_blocks_to_prompt_text(blocks: Sequence[ResourceBlock]) -> str:
|
||||
parts = []
|
||||
|
||||
for block in blocks:
|
||||
block_text = _format_content_block(block)
|
||||
if block_text is not None:
|
||||
parts.append(block_text)
|
||||
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _format_content_block(block: ResourceBlock) -> str | None:
|
||||
match block.get("type"):
|
||||
case "text":
|
||||
return block.get("text") or ""
|
||||
|
||||
case "resource":
|
||||
block_content = block.get("text") or ""
|
||||
fence = "```"
|
||||
return f"{block.get('uri')}\n{fence}\n{block_content}\n{fence}"
|
||||
|
||||
case "resource_link":
|
||||
fields = {
|
||||
"uri": block.get("uri"),
|
||||
"name": block.get("name"),
|
||||
"title": block.get("title"),
|
||||
"description": block.get("description"),
|
||||
"mimeType": block.get("mimeType"),
|
||||
"size": block.get("size"),
|
||||
}
|
||||
parts = [
|
||||
f"{k}: {v}"
|
||||
for k, v in fields.items()
|
||||
if v is not None and (v or isinstance(v, (int, float)))
|
||||
]
|
||||
return "\n".join(parts)
|
||||
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
BINARY_MIME_PREFIXES = (
|
||||
"audio/",
|
||||
"image/",
|
||||
"video/",
|
||||
"application/zip",
|
||||
"application/x-zip-compressed",
|
||||
)
|
||||
|
||||
|
||||
def _is_probably_text(path: PathResource, data: bytes) -> bool:
|
||||
mime_guess, _ = mimetypes.guess_type(path.path.name)
|
||||
if mime_guess and mime_guess.startswith(BINARY_MIME_PREFIXES):
|
||||
return False
|
||||
|
||||
if not data:
|
||||
return True
|
||||
if b"\x00" in data:
|
||||
return False
|
||||
|
||||
DEL_CODE = 127
|
||||
NON_PRINTABLE_MAX_PROPORTION = 0.1
|
||||
NON_PRINTABLE_MAX_CODE = 31
|
||||
NON_PRINTABLE_EXCEPTIONS = [9, 10, 11, 12]
|
||||
non_text = sum(
|
||||
1
|
||||
for b in data
|
||||
if b <= NON_PRINTABLE_MAX_CODE
|
||||
and b not in NON_PRINTABLE_EXCEPTIONS
|
||||
or b == DEL_CODE
|
||||
)
|
||||
return (non_text / len(data)) < NON_PRINTABLE_MAX_PROPORTION
|
||||
566
vibe/core/config.py
Normal file
566
vibe/core/config.py
Normal file
|
|
@ -0,0 +1,566 @@
|
|||
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.prompts import SystemPrompt
|
||||
from vibe.core.tools.base import BaseToolConfig
|
||||
|
||||
|
||||
def get_vibe_home() -> Path:
|
||||
if vibe_home := os.getenv("VIBE_HOME"):
|
||||
return Path(vibe_home).expanduser().resolve()
|
||||
return Path.home() / ".vibe"
|
||||
|
||||
|
||||
GLOBAL_CONFIG_DIR = get_vibe_home()
|
||||
GLOBAL_CONFIG_FILE = GLOBAL_CONFIG_DIR / "config.toml"
|
||||
GLOBAL_ENV_FILE = GLOBAL_CONFIG_DIR / ".env"
|
||||
|
||||
|
||||
def resolve_config_file() -> Path:
|
||||
for directory in (cwd := Path.cwd(), *cwd.parents):
|
||||
if (candidate := directory / ".vibe" / "config.toml").is_file():
|
||||
return candidate
|
||||
return GLOBAL_CONFIG_FILE
|
||||
|
||||
|
||||
def load_api_keys_from_env() -> None:
|
||||
if GLOBAL_ENV_FILE.is_file():
|
||||
env_vars = dotenv_values(GLOBAL_ENV_FILE)
|
||||
for key, value in env_vars.items():
|
||||
if value:
|
||||
os.environ.setdefault(key, value)
|
||||
|
||||
|
||||
CONFIG_FILE = resolve_config_file()
|
||||
CONFIG_DIR = CONFIG_FILE.parent
|
||||
AGENT_DIR = CONFIG_DIR / "agents"
|
||||
PROMPT_DIR = CONFIG_DIR / "prompts"
|
||||
INSTRUCTIONS_FILE = CONFIG_DIR / "instructions.md"
|
||||
HISTORY_FILE = CONFIG_DIR / "vibehistory"
|
||||
PROJECT_DOC_FILENAMES = ["AGENTS.md", "VIBE.md", ".vibe.md"]
|
||||
|
||||
|
||||
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
|
||||
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(get_vibe_home() / "logs" / "session")
|
||||
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
|
||||
|
||||
|
||||
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"
|
||||
vim_keybindings: bool = False
|
||||
disable_welcome_banner_animation: bool = False
|
||||
displayed_workdir: str = ""
|
||||
auto_compact_threshold: int = 100_000
|
||||
context_warnings: bool = False
|
||||
textual_theme: str = "textual-dark"
|
||||
instructions: str = ""
|
||||
workdir: Path | None = Field(default=None, exclude=True)
|
||||
system_prompt_id: str = "cli"
|
||||
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[str] = 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."
|
||||
),
|
||||
)
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="VIBE_", case_sensitive=False, extra="forbid"
|
||||
)
|
||||
|
||||
@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 / self.system_prompt_id).with_suffix(".md")
|
||||
if not custom_sp_path.is_file():
|
||||
raise MissingPromptFileError(self.system_prompt_id, str(PROMPT_DIR))
|
||||
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("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.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.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 / 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}"
|
||||
)
|
||||
|
||||
@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
|
||||
245
vibe/core/interaction_logger.py
Normal file
245
vibe/core/interaction_logger.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
import getpass
|
||||
import json
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import aiofiles
|
||||
|
||||
from vibe.core.llm.format import get_active_tool_classes
|
||||
from vibe.core.types import AgentStats, LLMMessage, SessionInfo, SessionMetadata
|
||||
from vibe.core.utils import is_windows
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import SessionLoggingConfig, VibeConfig
|
||||
from vibe.core.tools.manager import ToolManager
|
||||
|
||||
|
||||
class InteractionLogger:
|
||||
def __init__(
|
||||
self,
|
||||
session_config: SessionLoggingConfig,
|
||||
session_id: str,
|
||||
auto_approve: bool = False,
|
||||
workdir: Path | None = None,
|
||||
) -> None:
|
||||
if workdir is None:
|
||||
workdir = Path.cwd()
|
||||
self.session_config = session_config
|
||||
self.enabled = session_config.enabled
|
||||
self.auto_approve = auto_approve
|
||||
self.workdir = workdir
|
||||
|
||||
if not self.enabled:
|
||||
self.save_dir: Path | None = None
|
||||
self.session_prefix: str | None = None
|
||||
self.session_id: str = "disabled"
|
||||
self.session_start_time: str = "N/A"
|
||||
self.filepath: Path | None = None
|
||||
self.session_metadata: SessionMetadata | None = None
|
||||
return
|
||||
|
||||
self.save_dir = Path(session_config.save_dir)
|
||||
self.session_prefix = session_config.session_prefix
|
||||
self.session_id = session_id
|
||||
self.session_start_time = datetime.now().isoformat()
|
||||
|
||||
self.save_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.filepath = self._get_save_filepath()
|
||||
self.session_metadata = self._initialize_session_metadata()
|
||||
|
||||
def _get_save_filepath(self) -> Path:
|
||||
if self.save_dir is None or self.session_prefix is None:
|
||||
raise RuntimeError("Cannot get filepath when logging is disabled")
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{self.session_prefix}_{timestamp}_{self.session_id[:8]}.json"
|
||||
return self.save_dir / filename
|
||||
|
||||
def _get_git_commit(self) -> str | None:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"],
|
||||
capture_output=True,
|
||||
cwd=self.workdir,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=5.0,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout:
|
||||
return result.stdout.strip()
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
return None
|
||||
|
||||
def _get_git_branch(self) -> str | None:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--abbrev-ref", "HEAD"],
|
||||
capture_output=True,
|
||||
cwd=self.workdir,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=5.0,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout:
|
||||
return result.stdout.strip()
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
return None
|
||||
|
||||
def _get_username(self) -> str:
|
||||
try:
|
||||
return getpass.getuser()
|
||||
except Exception:
|
||||
return "unknown"
|
||||
|
||||
def _initialize_session_metadata(self) -> SessionMetadata:
|
||||
git_commit = self._get_git_commit()
|
||||
git_branch = self._get_git_branch()
|
||||
user_name = self._get_username()
|
||||
|
||||
return SessionMetadata(
|
||||
session_id=self.session_id,
|
||||
start_time=self.session_start_time,
|
||||
end_time=None,
|
||||
git_commit=git_commit,
|
||||
git_branch=git_branch,
|
||||
auto_approve=self.auto_approve,
|
||||
username=user_name,
|
||||
environment={"working_directory": str(self.workdir)},
|
||||
)
|
||||
|
||||
async def save_interaction(
|
||||
self,
|
||||
messages: list[LLMMessage],
|
||||
stats: AgentStats,
|
||||
config: VibeConfig,
|
||||
tool_manager: ToolManager,
|
||||
) -> str | None:
|
||||
if not self.enabled or self.filepath is None:
|
||||
return None
|
||||
|
||||
if self.session_metadata is None:
|
||||
return None
|
||||
|
||||
active_tools = get_active_tool_classes(tool_manager, config)
|
||||
|
||||
tools_available = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_class.get_name(),
|
||||
"description": tool_class.description,
|
||||
"parameters": tool_class.get_parameters(),
|
||||
},
|
||||
}
|
||||
for tool_class in active_tools
|
||||
]
|
||||
|
||||
interaction_data = {
|
||||
"metadata": {
|
||||
**self.session_metadata.model_dump(),
|
||||
"end_time": datetime.now().isoformat(),
|
||||
"stats": stats.model_dump(),
|
||||
"total_messages": len(messages),
|
||||
"tools_available": tools_available,
|
||||
"agent_config": config.model_dump(mode="json"),
|
||||
},
|
||||
"messages": [m.model_dump(exclude_none=True) for m in messages],
|
||||
}
|
||||
|
||||
try:
|
||||
json_content = json.dumps(interaction_data, indent=2, ensure_ascii=False)
|
||||
|
||||
async with aiofiles.open(self.filepath, "w", encoding="utf-8") as f:
|
||||
await f.write(json_content)
|
||||
|
||||
return str(self.filepath)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def reset_session(self, session_id: str) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
self.session_id = session_id
|
||||
self.session_start_time = datetime.now().isoformat()
|
||||
self.filepath = self._get_save_filepath()
|
||||
self.session_metadata = self._initialize_session_metadata()
|
||||
|
||||
def get_session_info(
|
||||
self, messages: list[dict[str, Any]], stats: AgentStats
|
||||
) -> SessionInfo:
|
||||
if not self.enabled or self.save_dir is None:
|
||||
return SessionInfo(
|
||||
session_id="disabled",
|
||||
start_time="N/A",
|
||||
message_count=len(messages),
|
||||
stats=stats,
|
||||
save_dir="N/A",
|
||||
)
|
||||
|
||||
return SessionInfo(
|
||||
session_id=self.session_id,
|
||||
start_time=self.session_start_time,
|
||||
message_count=len(messages),
|
||||
stats=stats,
|
||||
save_dir=str(self.save_dir),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def find_latest_session(config: SessionLoggingConfig) -> Path | None:
|
||||
save_dir = Path(config.save_dir)
|
||||
if not save_dir.exists():
|
||||
return None
|
||||
|
||||
pattern = f"{config.session_prefix}_*.json"
|
||||
session_files = list(save_dir.glob(pattern))
|
||||
|
||||
if not session_files:
|
||||
return None
|
||||
|
||||
return max(session_files, key=lambda p: p.stat().st_mtime)
|
||||
|
||||
@staticmethod
|
||||
def find_session_by_id(
|
||||
session_id: str, config: SessionLoggingConfig
|
||||
) -> Path | None:
|
||||
save_dir = Path(config.save_dir)
|
||||
if not save_dir.exists():
|
||||
return None
|
||||
|
||||
# If it's a full UUID, extract the short form (first 8 chars)
|
||||
short_id = session_id.split("-")[0] if "-" in session_id else session_id
|
||||
|
||||
# Try exact match first, then partial
|
||||
patterns = [
|
||||
f"{config.session_prefix}_*_{short_id}.json", # Exact short UUID
|
||||
f"{config.session_prefix}_*_{short_id}*.json", # Partial UUID
|
||||
]
|
||||
|
||||
for pattern in patterns:
|
||||
matches = list(save_dir.glob(pattern))
|
||||
if matches:
|
||||
return (
|
||||
max(matches, key=lambda p: p.stat().st_mtime)
|
||||
if len(matches) > 1
|
||||
else matches[0]
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def load_session(filepath: Path) -> tuple[list[LLMMessage], dict[str, Any]]:
|
||||
with filepath.open("r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
data = json.loads(content)
|
||||
messages = [LLMMessage.model_validate(msg) for msg in data.get("messages", [])]
|
||||
metadata = data.get("metadata", {})
|
||||
|
||||
return messages, metadata
|
||||
0
vibe/core/llm/__init__.py
Normal file
0
vibe/core/llm/__init__.py
Normal file
0
vibe/core/llm/backend/__init__.py
Normal file
0
vibe/core/llm/backend/__init__.py
Normal file
7
vibe/core/llm/backend/factory.py
Normal file
7
vibe/core/llm/backend/factory.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.core.config import Backend
|
||||
from vibe.core.llm.backend.generic import GenericBackend
|
||||
from vibe.core.llm.backend.mistral import MistralBackend
|
||||
|
||||
BACKEND_FACTORY = {Backend.MISTRAL: MistralBackend, Backend.GENERIC: GenericBackend}
|
||||
409
vibe/core/llm/backend/generic.py
Normal file
409
vibe/core/llm/backend/generic.py
Normal file
|
|
@ -0,0 +1,409 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
import json
|
||||
import os
|
||||
import types
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, Protocol, TypeVar
|
||||
|
||||
import httpx
|
||||
|
||||
from vibe.core.llm.exceptions import BackendErrorBuilder
|
||||
from vibe.core.types import (
|
||||
AvailableTool,
|
||||
LLMChunk,
|
||||
LLMMessage,
|
||||
LLMUsage,
|
||||
Role,
|
||||
StrToolChoice,
|
||||
)
|
||||
from vibe.core.utils import async_generator_retry, async_retry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import ModelConfig, ProviderConfig
|
||||
|
||||
|
||||
class PreparedRequest(NamedTuple):
|
||||
endpoint: str
|
||||
headers: dict[str, str]
|
||||
body: bytes
|
||||
|
||||
|
||||
class APIAdapter(Protocol):
|
||||
endpoint: ClassVar[str]
|
||||
|
||||
def prepare_request(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float,
|
||||
tools: list[AvailableTool] | None,
|
||||
max_tokens: int | None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
enable_streaming: bool,
|
||||
provider: ProviderConfig,
|
||||
api_key: str | None = None,
|
||||
) -> PreparedRequest: ...
|
||||
|
||||
def parse_response(self, data: dict[str, Any]) -> LLMChunk: ...
|
||||
|
||||
|
||||
BACKEND_ADAPTERS: dict[str, APIAdapter] = {}
|
||||
|
||||
T = TypeVar("T", bound=APIAdapter)
|
||||
|
||||
|
||||
def register_adapter(
|
||||
adapters: dict[str, APIAdapter], name: str
|
||||
) -> Callable[[type[T]], type[T]]:
|
||||
|
||||
def decorator(cls: type[T]) -> type[T]:
|
||||
adapters[name] = cls()
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@register_adapter(BACKEND_ADAPTERS, "openai")
|
||||
class OpenAIAdapter(APIAdapter):
|
||||
endpoint: ClassVar[str] = "/chat/completions"
|
||||
|
||||
def build_payload(
|
||||
self,
|
||||
model_name: str,
|
||||
converted_messages: list[dict[str, Any]],
|
||||
temperature: float,
|
||||
tools: list[AvailableTool] | None,
|
||||
max_tokens: int | None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
) -> dict[str, Any]:
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"messages": converted_messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
|
||||
if tools:
|
||||
payload["tools"] = [tool.model_dump(exclude_none=True) for tool in tools]
|
||||
if tool_choice:
|
||||
payload["tool_choice"] = (
|
||||
tool_choice
|
||||
if isinstance(tool_choice, str)
|
||||
else tool_choice.model_dump()
|
||||
)
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
return payload
|
||||
|
||||
def build_headers(self, api_key: str | None = None) -> dict[str, str]:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
def prepare_request(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float,
|
||||
tools: list[AvailableTool] | None,
|
||||
max_tokens: int | None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
enable_streaming: bool,
|
||||
provider: ProviderConfig,
|
||||
api_key: str | None = None,
|
||||
) -> PreparedRequest:
|
||||
converted_messages = [msg.model_dump(exclude_none=True) for msg in messages]
|
||||
|
||||
payload = self.build_payload(
|
||||
model_name, converted_messages, temperature, tools, max_tokens, tool_choice
|
||||
)
|
||||
|
||||
headers = self.build_headers(api_key)
|
||||
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
|
||||
return PreparedRequest(self.endpoint, headers, body)
|
||||
|
||||
def parse_response(self, data: dict[str, Any]) -> LLMChunk:
|
||||
if data.get("choices"):
|
||||
if "message" in data["choices"][0]:
|
||||
message = LLMMessage.model_validate(data["choices"][0]["message"])
|
||||
elif "delta" in data["choices"][0]:
|
||||
message = LLMMessage.model_validate(data["choices"][0]["delta"])
|
||||
else:
|
||||
raise ValueError("Invalid response data")
|
||||
finish_reason = data["choices"][0]["finish_reason"]
|
||||
|
||||
elif "message" in data:
|
||||
message = LLMMessage.model_validate(data["message"])
|
||||
finish_reason = data["finish_reason"]
|
||||
elif "delta" in data:
|
||||
message = LLMMessage.model_validate(data["delta"])
|
||||
finish_reason = None
|
||||
else:
|
||||
message = LLMMessage(role=Role.assistant, content="")
|
||||
finish_reason = None
|
||||
|
||||
usage_data = data.get("usage") or {}
|
||||
usage = LLMUsage(
|
||||
prompt_tokens=usage_data.get("prompt_tokens", 0),
|
||||
completion_tokens=usage_data.get("completion_tokens", 0),
|
||||
)
|
||||
|
||||
return LLMChunk(message=message, usage=usage, finish_reason=finish_reason)
|
||||
|
||||
|
||||
class GenericBackend:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: httpx.AsyncClient | None = None,
|
||||
provider: ProviderConfig,
|
||||
timeout: float = 720.0,
|
||||
) -> None:
|
||||
"""Initialize the backend.
|
||||
|
||||
Args:
|
||||
client: Optional httpx client to use. If not provided, one will be created.
|
||||
"""
|
||||
self._client = client
|
||||
self._owns_client = client is None
|
||||
self._provider = provider
|
||||
self._timeout = timeout
|
||||
|
||||
async def __aenter__(self) -> GenericBackend:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._timeout),
|
||||
limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
|
||||
)
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: types.TracebackType | None,
|
||||
) -> None:
|
||||
if self._owns_client and self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
def _get_client(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._timeout),
|
||||
limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
|
||||
)
|
||||
self._owns_client = True
|
||||
return self._client
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
*,
|
||||
model: ModelConfig,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float = 0.2,
|
||||
tools: list[AvailableTool] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> LLMChunk:
|
||||
api_key = (
|
||||
os.getenv(self._provider.api_key_env_var)
|
||||
if self._provider.api_key_env_var
|
||||
else None
|
||||
)
|
||||
|
||||
api_style = getattr(self._provider, "api_style", "openai")
|
||||
adapter = BACKEND_ADAPTERS[api_style]
|
||||
|
||||
endpoint, headers, body = adapter.prepare_request(
|
||||
model_name=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
tool_choice=tool_choice,
|
||||
enable_streaming=False,
|
||||
provider=self._provider,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
url = f"{self._provider.api_base}{endpoint}"
|
||||
|
||||
try:
|
||||
res_data, _ = await self._make_request(url, body, headers)
|
||||
return adapter.parse_response(res_data)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise BackendErrorBuilder.build_http_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=url,
|
||||
response=e.response,
|
||||
headers=dict(e.response.headers.items()),
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
except httpx.RequestError as e:
|
||||
raise BackendErrorBuilder.build_request_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=url,
|
||||
error=e,
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
|
||||
async def complete_streaming(
|
||||
self,
|
||||
*,
|
||||
model: ModelConfig,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float = 0.2,
|
||||
tools: list[AvailableTool] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> AsyncGenerator[LLMChunk, None]:
|
||||
api_key = (
|
||||
os.getenv(self._provider.api_key_env_var)
|
||||
if self._provider.api_key_env_var
|
||||
else None
|
||||
)
|
||||
|
||||
api_style = getattr(self._provider, "api_style", "openai")
|
||||
adapter = BACKEND_ADAPTERS[api_style]
|
||||
|
||||
endpoint, headers, body = adapter.prepare_request(
|
||||
model_name=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
tool_choice=tool_choice,
|
||||
enable_streaming=True,
|
||||
provider=self._provider,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
url = f"{self._provider.api_base}{endpoint}"
|
||||
|
||||
try:
|
||||
async for res_data in self._make_streaming_request(url, body, headers):
|
||||
yield adapter.parse_response(res_data)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise BackendErrorBuilder.build_http_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=url,
|
||||
response=e.response,
|
||||
headers=dict(e.response.headers.items()),
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
except httpx.RequestError as e:
|
||||
raise BackendErrorBuilder.build_request_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=url,
|
||||
error=e,
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
|
||||
class HTTPResponse(NamedTuple):
|
||||
data: dict[str, Any]
|
||||
headers: dict[str, str]
|
||||
|
||||
@async_retry(tries=3)
|
||||
async def _make_request(
|
||||
self, url: str, data: bytes, headers: dict[str, str]
|
||||
) -> HTTPResponse:
|
||||
client = self._get_client()
|
||||
response = await client.post(url, content=data, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
response_headers = dict(response.headers.items())
|
||||
response_body = response.json()
|
||||
return self.HTTPResponse(response_body, response_headers)
|
||||
|
||||
@async_generator_retry(tries=3)
|
||||
async def _make_streaming_request(
|
||||
self, url: str, data: bytes, headers: dict[str, str]
|
||||
) -> AsyncGenerator[dict[str, Any]]:
|
||||
client = self._get_client()
|
||||
async with client.stream(
|
||||
method="POST", url=url, content=data, headers=headers
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
async for line in response.aiter_lines():
|
||||
if line.strip() == "":
|
||||
continue
|
||||
|
||||
DELIM_CHAR = ":"
|
||||
assert f"{DELIM_CHAR} " in line, "line should look like `key: value`"
|
||||
delim_index = line.find(DELIM_CHAR)
|
||||
key = line[0:delim_index]
|
||||
value = line[delim_index + 2 :]
|
||||
|
||||
if key != "data":
|
||||
# This might be the case with openrouter, so we just ignore it
|
||||
continue
|
||||
if value == "[DONE]":
|
||||
return
|
||||
yield json.loads(value.strip())
|
||||
|
||||
async def count_tokens(
|
||||
self,
|
||||
*,
|
||||
model: ModelConfig,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float = 0.0,
|
||||
tools: list[AvailableTool] | None = None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> int:
|
||||
probe_messages = list(messages)
|
||||
if not probe_messages or probe_messages[-1].role != Role.user:
|
||||
probe_messages.append(LLMMessage(role=Role.user, content=""))
|
||||
|
||||
result = await self.complete(
|
||||
model=model,
|
||||
messages=probe_messages,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
max_tokens=16, # Minimal amount for openrouter with openai models
|
||||
tool_choice=tool_choice,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
assert result.usage is not None, (
|
||||
"Usage should be present in non-streaming completions"
|
||||
)
|
||||
|
||||
return result.usage.prompt_tokens
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._owns_client and self._client:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
332
vibe/core/llm/backend/mistral.py
Normal file
332
vibe/core/llm/backend/mistral.py
Normal file
|
|
@ -0,0 +1,332 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import types
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import httpx
|
||||
import mistralai
|
||||
|
||||
from vibe.core.llm.exceptions import BackendErrorBuilder
|
||||
from vibe.core.types import (
|
||||
AvailableTool,
|
||||
Content,
|
||||
FunctionCall,
|
||||
LLMChunk,
|
||||
LLMMessage,
|
||||
LLMUsage,
|
||||
Role,
|
||||
StrToolChoice,
|
||||
ToolCall,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import ModelConfig, ProviderConfig
|
||||
|
||||
|
||||
class MistralMapper:
|
||||
def prepare_message(self, msg: LLMMessage) -> mistralai.Messages:
|
||||
match msg.role:
|
||||
case Role.system:
|
||||
return mistralai.SystemMessage(role="system", content=msg.content or "")
|
||||
case Role.user:
|
||||
return mistralai.UserMessage(role="user", content=msg.content)
|
||||
case Role.assistant:
|
||||
return mistralai.AssistantMessage(
|
||||
role="assistant",
|
||||
content=msg.content,
|
||||
tool_calls=[
|
||||
mistralai.ToolCall(
|
||||
function=mistralai.FunctionCall(
|
||||
name=tc.function.name or "",
|
||||
arguments=tc.function.arguments or "",
|
||||
),
|
||||
id=tc.id,
|
||||
type=tc.type,
|
||||
index=tc.index,
|
||||
)
|
||||
for tc in msg.tool_calls or []
|
||||
],
|
||||
)
|
||||
case Role.tool:
|
||||
return mistralai.ToolMessage(
|
||||
role="tool",
|
||||
content=msg.content,
|
||||
tool_call_id=msg.tool_call_id,
|
||||
name=msg.name,
|
||||
)
|
||||
|
||||
def prepare_tool(self, tool: AvailableTool) -> mistralai.Tool:
|
||||
return mistralai.Tool(
|
||||
type="function",
|
||||
function=mistralai.Function(
|
||||
name=tool.function.name,
|
||||
description=tool.function.description,
|
||||
parameters=tool.function.parameters,
|
||||
),
|
||||
)
|
||||
|
||||
def prepare_tool_choice(
|
||||
self, tool_choice: StrToolChoice | AvailableTool
|
||||
) -> mistralai.ChatCompletionStreamRequestToolChoice:
|
||||
if isinstance(tool_choice, str):
|
||||
return cast(mistralai.ToolChoiceEnum, tool_choice)
|
||||
|
||||
return mistralai.ToolChoice(
|
||||
type="function",
|
||||
function=mistralai.FunctionName(name=tool_choice.function.name),
|
||||
)
|
||||
|
||||
def parse_content(self, content: mistralai.AssistantMessageContent) -> Content:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
|
||||
concat_content = ""
|
||||
for chunk in content:
|
||||
if isinstance(chunk, mistralai.FileChunk):
|
||||
continue
|
||||
match chunk.type:
|
||||
case "text":
|
||||
concat_content += chunk.text
|
||||
case _:
|
||||
pass
|
||||
return concat_content
|
||||
|
||||
def parse_tool_calls(self, tool_calls: list[mistralai.ToolCall]) -> list[ToolCall]:
|
||||
return [
|
||||
ToolCall(
|
||||
id=tool_call.id,
|
||||
function=FunctionCall(
|
||||
name=tool_call.function.name,
|
||||
arguments=tool_call.function.arguments
|
||||
if isinstance(tool_call.function.arguments, str)
|
||||
else json.dumps(tool_call.function.arguments),
|
||||
),
|
||||
index=tool_call.index,
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
|
||||
|
||||
class MistralBackend:
|
||||
def __init__(self, provider: ProviderConfig, timeout: float = 720.0) -> None:
|
||||
self._client: mistralai.Mistral | None = None
|
||||
self._provider = provider
|
||||
self._mapper = MistralMapper()
|
||||
self._api_key = (
|
||||
os.getenv(self._provider.api_key_env_var)
|
||||
if self._provider.api_key_env_var
|
||||
else None
|
||||
)
|
||||
|
||||
# Mistral SDK takes server URL without api version as input
|
||||
url_pattern = r"(https?://[^/]+)(/v.*)"
|
||||
match = re.match(url_pattern, self._provider.api_base)
|
||||
if not match:
|
||||
raise ValueError(
|
||||
f"Invalid API base URL: {self._provider.api_base}. "
|
||||
"Expected format: <server_url>/v<api_version>"
|
||||
)
|
||||
self._server_url = match.group(1)
|
||||
self._timeout = timeout
|
||||
|
||||
async def __aenter__(self) -> MistralBackend:
|
||||
self._client = mistralai.Mistral(
|
||||
api_key=self._api_key,
|
||||
server_url=self._server_url,
|
||||
timeout_ms=int(self._timeout * 1000),
|
||||
)
|
||||
await self._client.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: types.TracebackType | None,
|
||||
) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.__aexit__(
|
||||
exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb
|
||||
)
|
||||
|
||||
def _get_client(self) -> mistralai.Mistral:
|
||||
if self._client is None:
|
||||
self._client = mistralai.Mistral(
|
||||
api_key=self._api_key, server_url=self._server_url
|
||||
)
|
||||
return self._client
|
||||
|
||||
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:
|
||||
try:
|
||||
response = await self._get_client().chat.complete_async(
|
||||
model=model.name,
|
||||
messages=[self._mapper.prepare_message(msg) for msg in messages],
|
||||
temperature=temperature,
|
||||
tools=[self._mapper.prepare_tool(tool) for tool in tools]
|
||||
if tools
|
||||
else None,
|
||||
max_tokens=max_tokens,
|
||||
tool_choice=self._mapper.prepare_tool_choice(tool_choice)
|
||||
if tool_choice
|
||||
else None,
|
||||
http_headers=extra_headers,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
return LLMChunk(
|
||||
message=LLMMessage(
|
||||
role=Role.assistant,
|
||||
content=self._mapper.parse_content(
|
||||
response.choices[0].message.content
|
||||
)
|
||||
if response.choices[0].message.content
|
||||
else "",
|
||||
tool_calls=self._mapper.parse_tool_calls(
|
||||
response.choices[0].message.tool_calls
|
||||
)
|
||||
if response.choices[0].message.tool_calls
|
||||
else None,
|
||||
),
|
||||
usage=LLMUsage(
|
||||
prompt_tokens=response.usage.prompt_tokens or 0,
|
||||
completion_tokens=response.usage.completion_tokens or 0,
|
||||
),
|
||||
finish_reason=response.choices[0].finish_reason,
|
||||
)
|
||||
|
||||
except mistralai.SDKError as e:
|
||||
raise BackendErrorBuilder.build_http_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=self._server_url,
|
||||
response=e.raw_response,
|
||||
headers=dict(e.raw_response.headers.items()),
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
except httpx.RequestError as e:
|
||||
raise BackendErrorBuilder.build_request_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=self._server_url,
|
||||
error=e,
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
|
||||
async 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]:
|
||||
try:
|
||||
async for chunk in await self._get_client().chat.stream_async(
|
||||
model=model.name,
|
||||
messages=[self._mapper.prepare_message(msg) for msg in messages],
|
||||
temperature=temperature,
|
||||
tools=[self._mapper.prepare_tool(tool) for tool in tools]
|
||||
if tools
|
||||
else None,
|
||||
max_tokens=max_tokens,
|
||||
tool_choice=self._mapper.prepare_tool_choice(tool_choice)
|
||||
if tool_choice
|
||||
else None,
|
||||
http_headers=extra_headers,
|
||||
):
|
||||
yield LLMChunk(
|
||||
message=LLMMessage(
|
||||
role=Role.assistant,
|
||||
content=self._mapper.parse_content(
|
||||
chunk.data.choices[0].delta.content
|
||||
)
|
||||
if chunk.data.choices[0].delta.content
|
||||
else "",
|
||||
tool_calls=self._mapper.parse_tool_calls(
|
||||
chunk.data.choices[0].delta.tool_calls
|
||||
)
|
||||
if chunk.data.choices[0].delta.tool_calls
|
||||
else None,
|
||||
),
|
||||
usage=LLMUsage(
|
||||
prompt_tokens=chunk.data.usage.prompt_tokens or 0
|
||||
if chunk.data.usage
|
||||
else 0,
|
||||
completion_tokens=chunk.data.usage.completion_tokens or 0
|
||||
if chunk.data.usage
|
||||
else 0,
|
||||
),
|
||||
finish_reason=chunk.data.choices[0].finish_reason,
|
||||
)
|
||||
|
||||
except mistralai.SDKError as e:
|
||||
raise BackendErrorBuilder.build_http_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=self._server_url,
|
||||
response=e.raw_response,
|
||||
headers=dict(e.raw_response.headers.items()),
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
except httpx.RequestError as e:
|
||||
raise BackendErrorBuilder.build_request_error(
|
||||
provider=self._provider.name,
|
||||
endpoint=self._server_url,
|
||||
error=e,
|
||||
model=model.name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
has_tools=bool(tools),
|
||||
tool_choice=tool_choice,
|
||||
) from e
|
||||
|
||||
async def count_tokens(
|
||||
self,
|
||||
*,
|
||||
model: ModelConfig,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float = 0.0,
|
||||
tools: list[AvailableTool] | None = None,
|
||||
tool_choice: StrToolChoice | AvailableTool | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> int:
|
||||
result = await self.complete(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
max_tokens=1,
|
||||
tool_choice=tool_choice,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
assert result.usage is not None, (
|
||||
"Usage should be present in non-streaming completions"
|
||||
)
|
||||
|
||||
return result.usage.prompt_tokens
|
||||
195
vibe/core/llm/exceptions.py
Normal file
195
vibe/core/llm/exceptions.py
Normal file
|
|
@ -0,0 +1,195 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from http import HTTPStatus
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from vibe.core.types import AvailableTool, LLMMessage, StrToolChoice
|
||||
|
||||
|
||||
class ErrorDetail(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
message: str | None = None
|
||||
|
||||
|
||||
class PayloadSummary(BaseModel):
|
||||
model: str
|
||||
message_count: int
|
||||
approx_chars: int
|
||||
temperature: float
|
||||
has_tools: bool
|
||||
tool_choice: StrToolChoice | AvailableTool | None
|
||||
|
||||
|
||||
class BackendError(RuntimeError):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
provider: str,
|
||||
endpoint: str,
|
||||
status: int | None,
|
||||
reason: str | None,
|
||||
headers: Mapping[str, str] | None,
|
||||
body_text: str | None,
|
||||
parsed_error: str | None,
|
||||
model: str,
|
||||
payload_summary: PayloadSummary,
|
||||
) -> None:
|
||||
self.provider = provider
|
||||
self.endpoint = endpoint
|
||||
self.status = status
|
||||
self.reason = reason
|
||||
self.headers = {k.lower(): v for k, v in (headers or {}).items()}
|
||||
self.body_text = body_text or ""
|
||||
self.parsed_error = parsed_error
|
||||
self.model = model
|
||||
self.payload_summary = payload_summary
|
||||
super().__init__(self._fmt())
|
||||
|
||||
def _fmt(self) -> str:
|
||||
if self.status == HTTPStatus.UNAUTHORIZED:
|
||||
return "Invalid API key. Please check your API key and try again."
|
||||
|
||||
if self.status == HTTPStatus.TOO_MANY_REQUESTS:
|
||||
return "Rate limit exceeded. Please wait a moment before trying again."
|
||||
|
||||
rid = self.headers.get("x-request-id") or self.headers.get("request-id")
|
||||
status_label = (
|
||||
f"{self.status} {HTTPStatus(self.status).phrase}" if self.status else "N/A"
|
||||
)
|
||||
parts = [
|
||||
f"LLM backend error [{self.provider}]",
|
||||
f" status: {status_label}",
|
||||
f" reason: {self.reason or 'N/A'}",
|
||||
f" request_id: {rid or 'N/A'}",
|
||||
f" endpoint: {self.endpoint}",
|
||||
f" model: {self.model}",
|
||||
f" provider_message: {self.parsed_error or 'N/A'}",
|
||||
f" body_excerpt: {self._excerpt(self.body_text)}",
|
||||
f" payload_summary: {self.payload_summary.model_dump_json(exclude_none=True)}",
|
||||
]
|
||||
return "\n".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def _excerpt(s: str, *, n: int = 400) -> str:
|
||||
s = s.strip().replace("\n", " ")
|
||||
return s[:n] + ("…" if len(s) > n else "")
|
||||
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
error: ErrorDetail | dict[str, Any] | None = None
|
||||
message: str | None = None
|
||||
detail: str | None = None
|
||||
|
||||
@property
|
||||
def primary_message(self) -> str | None:
|
||||
if e := self.error:
|
||||
match e:
|
||||
case {"message": str(m)}:
|
||||
return m
|
||||
case {"type": str(t)}:
|
||||
return f"Error: {t}"
|
||||
case ErrorDetail(message=str(m)):
|
||||
return m
|
||||
if m := self.message:
|
||||
return m
|
||||
if d := self.detail:
|
||||
return d
|
||||
return None
|
||||
|
||||
|
||||
class BackendErrorBuilder:
|
||||
@classmethod
|
||||
def build_http_error(
|
||||
cls,
|
||||
*,
|
||||
provider: str,
|
||||
endpoint: str,
|
||||
response: httpx.Response,
|
||||
headers: Mapping[str, str] | None,
|
||||
model: str,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float,
|
||||
has_tools: bool,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
) -> BackendError:
|
||||
try:
|
||||
body_text = response.text
|
||||
except Exception: # On streaming responses, we can't read the body
|
||||
body_text = None
|
||||
|
||||
return BackendError(
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
status=response.status_code,
|
||||
reason=response.reason_phrase,
|
||||
headers=headers or {},
|
||||
body_text=body_text,
|
||||
parsed_error=cls._parse_provider_error(body_text),
|
||||
model=model,
|
||||
payload_summary=cls._payload_summary(
|
||||
model, messages, temperature, has_tools, tool_choice
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_request_error(
|
||||
cls,
|
||||
*,
|
||||
provider: str,
|
||||
endpoint: str,
|
||||
error: httpx.RequestError,
|
||||
model: str,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float,
|
||||
has_tools: bool,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
) -> BackendError:
|
||||
return BackendError(
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
status=None,
|
||||
reason=str(error) or repr(error),
|
||||
headers={},
|
||||
body_text=None,
|
||||
parsed_error="Network error",
|
||||
model=model,
|
||||
payload_summary=cls._payload_summary(
|
||||
model, messages, temperature, has_tools, tool_choice
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_provider_error(body_text: str | None) -> str | None:
|
||||
if not body_text:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(body_text)
|
||||
error_model = ErrorResponse.model_validate(data)
|
||||
return error_model.primary_message
|
||||
except (json.JSONDecodeError, ValidationError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _payload_summary(
|
||||
model_name: str,
|
||||
messages: list[LLMMessage],
|
||||
temperature: float,
|
||||
has_tools: bool,
|
||||
tool_choice: StrToolChoice | AvailableTool | None,
|
||||
) -> PayloadSummary:
|
||||
total_chars = sum(len(m.content or "") for m in messages)
|
||||
return PayloadSummary(
|
||||
model=model_name,
|
||||
message_count=len(messages),
|
||||
approx_chars=total_chars,
|
||||
temperature=temperature,
|
||||
has_tools=has_tools,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
270
vibe/core/llm/format.py
Normal file
270
vibe/core/llm/format.py
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from fnmatch import fnmatch
|
||||
from functools import lru_cache
|
||||
import json
|
||||
import re
|
||||
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.config import VibeConfig
|
||||
from vibe.core.tools.manager import ToolManager
|
||||
|
||||
|
||||
def _is_regex_hint(pattern: str) -> bool:
|
||||
"""Heuristically detect whether a pattern looks like a regex.
|
||||
|
||||
- Explicit regex: starts with 're:'
|
||||
- Heuristic regex: contains common regex metachars or '.*'
|
||||
"""
|
||||
if pattern.startswith("re:"):
|
||||
return True
|
||||
return bool(re.search(r"[().+|^$]", pattern) or ".*" in pattern)
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _compile_icase(expr: str) -> re.Pattern | None:
|
||||
try:
|
||||
return re.compile(expr, re.IGNORECASE)
|
||||
except re.error:
|
||||
return None
|
||||
|
||||
|
||||
def _regex_match_icase(expr: str, s: str) -> bool:
|
||||
rx = _compile_icase(expr)
|
||||
return rx is not None and rx.fullmatch(s) is not None
|
||||
|
||||
|
||||
def _name_matches(name: str, patterns: list[str]) -> bool:
|
||||
"""Check if a tool name matches any of the provided patterns.
|
||||
|
||||
Supports three forms (case-insensitive):
|
||||
- Exact names (no wildcards/regex tokens)
|
||||
- Glob wildcards using fnmatch (e.g., 'serena_*')
|
||||
- Regex when prefixed with 're:'
|
||||
or when the pattern looks regex-y (e.g., 'serena.*')
|
||||
"""
|
||||
n = name.lower()
|
||||
for raw in patterns:
|
||||
if not (p := (raw or "").strip()):
|
||||
continue
|
||||
|
||||
match p:
|
||||
case _ if p.startswith("re:"):
|
||||
if _regex_match_icase(p.removeprefix("re:"), name):
|
||||
return True
|
||||
case _ if _is_regex_hint(p):
|
||||
if _regex_match_icase(p, name):
|
||||
return True
|
||||
case _:
|
||||
if fnmatch(n, p.lower()):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def get_active_tool_classes(
|
||||
tool_manager: ToolManager, config: VibeConfig
|
||||
) -> list[type[BaseTool]]:
|
||||
"""Returns a list of active tool classes based on the configuration.
|
||||
|
||||
Args:
|
||||
tool_manager: ToolManager instance with discovered tools
|
||||
config: VibeConfig with enabled_tools/disabled_tools settings
|
||||
"""
|
||||
all_tools = list(tool_manager.available_tools().values())
|
||||
|
||||
if config.enabled_tools:
|
||||
return [
|
||||
tool_class
|
||||
for tool_class in all_tools
|
||||
if _name_matches(tool_class.get_name(), config.enabled_tools)
|
||||
]
|
||||
|
||||
if config.disabled_tools:
|
||||
return [
|
||||
tool_class
|
||||
for tool_class in all_tools
|
||||
if not _name_matches(tool_class.get_name(), config.disabled_tools)
|
||||
]
|
||||
|
||||
return all_tools
|
||||
|
||||
|
||||
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, config: VibeConfig
|
||||
) -> list[AvailableTool]:
|
||||
active_tools = get_active_tool_classes(tool_manager, config)
|
||||
|
||||
return [
|
||||
AvailableTool(
|
||||
function=AvailableFunction(
|
||||
name=tool_class.get_name(),
|
||||
description=tool_class.description,
|
||||
parameters=tool_class.get_parameters(),
|
||||
)
|
||||
)
|
||||
for tool_class in active_tools
|
||||
]
|
||||
|
||||
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}
|
||||
|
||||
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, config: VibeConfig
|
||||
) -> ResolvedMessage:
|
||||
resolved_calls = []
|
||||
failed_calls = []
|
||||
|
||||
active_tools = {
|
||||
tool_class.get_name(): tool_class
|
||||
for tool_class in get_active_tool_classes(tool_manager, config)
|
||||
}
|
||||
|
||||
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,
|
||||
)
|
||||
120
vibe/core/llm/types.py
Normal file
120
vibe/core/llm/types.py
Normal file
|
|
@ -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 Agent 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
|
||||
"""
|
||||
...
|
||||
191
vibe/core/middleware.py
Normal file
191
vibe/core/middleware.py
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum, auto
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
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: ...
|
||||
|
||||
async def after_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()
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
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()
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
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()
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
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()
|
||||
|
||||
async def after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
return MiddlewareResult()
|
||||
|
||||
def reset(self, reset_reason: ResetReason = ResetReason.STOP) -> None:
|
||||
self.has_warned = 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()
|
||||
|
||||
async def run_after_turn(self, context: ConversationContext) -> MiddlewareResult:
|
||||
messages_to_inject = []
|
||||
|
||||
for mw in self.middlewares:
|
||||
result = await mw.after_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()
|
||||
85
vibe/core/output_formatters.py
Normal file
85
vibe/core/output_formatters.py
Normal file
|
|
@ -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)
|
||||
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)
|
||||
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)
|
||||
64
vibe/core/programmatic.py
Normal file
64
vibe/core/programmatic.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from vibe.core.agent import Agent
|
||||
from vibe.core.config import VibeConfig
|
||||
from vibe.core.output_formatters import create_formatter
|
||||
from vibe.core.types import AssistantEvent, LLMMessage, OutputFormat, Role
|
||||
from vibe.core.utils import ConversationLimitException, logger
|
||||
|
||||
|
||||
def run_programmatic(
|
||||
config: VibeConfig,
|
||||
prompt: str,
|
||||
max_turns: int | None = None,
|
||||
max_price: float | None = None,
|
||||
output_format: OutputFormat = OutputFormat.TEXT,
|
||||
previous_messages: list[LLMMessage] | None = None,
|
||||
auto_approve: bool = True,
|
||||
) -> str | None:
|
||||
"""Run in programmatic mode: execute prompt and return the assistant response.
|
||||
|
||||
Args:
|
||||
config: Configuration for the Vibe agent
|
||||
prompt: The user prompt to process
|
||||
max_turns: Maximum number of assistant turns (LLM calls) to allow
|
||||
max_price: Maximum cost in dollars before stopping
|
||||
output_format: Format for the output
|
||||
previous_messages: Optional messages from a previous session to continue
|
||||
auto_approve: Whether to automatically approve tool execution
|
||||
|
||||
Returns:
|
||||
The final assistant response text, or None if no response
|
||||
"""
|
||||
formatter = create_formatter(output_format)
|
||||
|
||||
agent = Agent(
|
||||
config,
|
||||
auto_approve=auto_approve,
|
||||
message_observer=formatter.on_message_added,
|
||||
max_turns=max_turns,
|
||||
max_price=max_price,
|
||||
enable_streaming=False,
|
||||
)
|
||||
logger.info("USER: %s", prompt)
|
||||
|
||||
async def _async_run() -> str | None:
|
||||
if previous_messages:
|
||||
non_system_messages = [
|
||||
msg for msg in previous_messages if not (msg.role == Role.system)
|
||||
]
|
||||
agent.messages.extend(non_system_messages)
|
||||
logger.info(
|
||||
"Loaded %d messages from previous session", len(non_system_messages)
|
||||
)
|
||||
|
||||
async for event in agent.act(prompt):
|
||||
formatter.on_event(event)
|
||||
if isinstance(event, AssistantEvent) and event.stopped_by_middleware:
|
||||
raise ConversationLimitException(event.content)
|
||||
|
||||
return formatter.finalize()
|
||||
|
||||
return asyncio.run(_async_run())
|
||||
31
vibe/core/prompts/__init__.py
Normal file
31
vibe/core/prompts/__init__.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum, auto
|
||||
from pathlib import Path
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
|
||||
_PROMPTS_DIR = VIBE_ROOT / "core" / "prompts"
|
||||
|
||||
|
||||
class Prompt(StrEnum):
|
||||
@property
|
||||
def path(self) -> Path:
|
||||
return (_PROMPTS_DIR / self.value).with_suffix(".md")
|
||||
|
||||
def read(self) -> str:
|
||||
return self.path.read_text(encoding="utf-8").strip()
|
||||
|
||||
|
||||
class SystemPrompt(Prompt):
|
||||
CLI = auto()
|
||||
TESTS = auto()
|
||||
|
||||
|
||||
class UtilityPrompt(Prompt):
|
||||
COMPACT = auto()
|
||||
DANGEROUS_DIRECTORY = auto()
|
||||
PROJECT_CONTEXT = auto()
|
||||
|
||||
|
||||
__all__ = ["SystemPrompt", "UtilityPrompt"]
|
||||
24
vibe/core/prompts/cli.md
Normal file
24
vibe/core/prompts/cli.md
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
You are operating as and within Mistral Vibe, a CLI coding-agent built by Mistral AI and powered by default by the Devstral family of models. It wraps Mistral's Devstral models to enable natural language interaction with a local codebase. Use the available tools when helpful.
|
||||
|
||||
You can:
|
||||
|
||||
- Receive user prompts, project context, and files.
|
||||
- Send responses and emit function calls (e.g., shell commands, code edits).
|
||||
- Apply patches, run commands, based on user approvals.
|
||||
|
||||
Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values; otherwise proceed with the tool calls. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.
|
||||
|
||||
Always try your hardest to use the tools to answer the user's request. If you can't use the tools, explain why and ask the user for more information.
|
||||
|
||||
Act as an agentic assistant, if a user asks for a long task, break it down and do it step by step.
|
||||
|
||||
When you want to commit changes, you will always use the 'git commit' bash command. It will always
|
||||
be suffixed with a line telling it was generated by Mistral Vibe with the appropriate co-authoring information.
|
||||
The format you will always uses is the following heredoc.
|
||||
|
||||
```bash
|
||||
git commit -m "<Commit message here>
|
||||
|
||||
Generated by Mistral Vibe.
|
||||
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>"
|
||||
```
|
||||
48
vibe/core/prompts/compact.md
Normal file
48
vibe/core/prompts/compact.md
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
Create a comprehensive summary of our entire conversation that will serve as complete context for continuing this work. Structure your summary to capture both the narrative flow and technical details necessary for seamless continuation.
|
||||
|
||||
Your summary must include these sections in order:
|
||||
|
||||
## 1. User's Primary Goals and Intent
|
||||
Capture ALL explicit requests and objectives stated by the user throughout the conversation, preserving their exact priorities and constraints.
|
||||
|
||||
## 2. Conversation Timeline and Progress
|
||||
Chronologically document the key phases of our work:
|
||||
- Initial requests and how they were addressed
|
||||
- Major decisions made and their rationale
|
||||
- Problems encountered and solutions applied
|
||||
- Current state of the work
|
||||
|
||||
## 3. Technical Context and Decisions
|
||||
- Technologies, frameworks, and tools being used
|
||||
- Architectural patterns and design decisions made
|
||||
- Key technical constraints or requirements identified
|
||||
- Important code patterns or conventions established
|
||||
|
||||
## 4. Files and Code Changes
|
||||
For each file created, modified, or examined:
|
||||
- Full file path/name
|
||||
- Purpose and importance of the file
|
||||
- Specific changes made (with key code snippets where critical)
|
||||
- Current state of the file
|
||||
|
||||
## 5. Active Work and Last Actions
|
||||
CRITICAL: Detail EXACTLY what was being worked on in the most recent exchanges:
|
||||
- The specific task or problem being addressed
|
||||
- Last completed action
|
||||
- Any partial work or mid-implementation state
|
||||
- Include relevant code snippets from the most recent work
|
||||
|
||||
## 6. Unresolved Issues and Pending Tasks
|
||||
- Any errors or issues still requiring attention
|
||||
- Tasks explicitly requested but not yet started
|
||||
- Decisions waiting for user input
|
||||
|
||||
## 7. Immediate Next Step
|
||||
State the SPECIFIC next action to take based on:
|
||||
- The user's most recent request
|
||||
- The current state of implementation
|
||||
- Any ongoing work that was interrupted
|
||||
|
||||
Important: Be precise with technical details, file names, and code. The next agent reading this should be able to continue exactly where we left off without asking clarifying questions. Include enough detail that no context is lost, but remain focused on actionable information.
|
||||
|
||||
Respond with ONLY the summary text following this structure - no additional commentary or meta-discussion.
|
||||
5
vibe/core/prompts/dangerous_directory.md
Normal file
5
vibe/core/prompts/dangerous_directory.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
directoryStructure: Project context scanning has been disabled because {reason}. This prevents permission dialogs and potential system slowdowns. Use the LS tool and other file tools to explore the project structure as needed.
|
||||
|
||||
Absolute path: {abs_path}
|
||||
|
||||
gitStatus: Use git tools to check repository status if needed.
|
||||
8
vibe/core/prompts/project_context.md
Normal file
8
vibe/core/prompts/project_context.md
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
directoryStructure: Below is a snapshot of this project's file structure at the start of the conversation. This snapshot will NOT update during the conversation. It skips over .gitignore patterns.{large_repo_warning}
|
||||
|
||||
{structure}
|
||||
|
||||
Absolute path: {abs_path}
|
||||
|
||||
gitStatus: This is the git status at the start of the conversation. Note that this status is a snapshot in time, and will not update during the conversation.
|
||||
{git_status}
|
||||
1
vibe/core/prompts/tests.md
Normal file
1
vibe/core/prompts/tests.md
Normal file
|
|
@ -0,0 +1 @@
|
|||
You are Vibe, a super useful programming assistant.
|
||||
404
vibe/core/system_prompt.py
Normal file
404
vibe/core/system_prompt.py
Normal file
|
|
@ -0,0 +1,404 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Generator
|
||||
import fnmatch
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from vibe.core.config import INSTRUCTIONS_FILE, PROJECT_DOC_FILENAMES
|
||||
from vibe.core.llm.format import get_active_tool_classes
|
||||
from vibe.core.prompts import UtilityPrompt
|
||||
from vibe.core.utils import is_dangerous_directory, is_windows
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import ProjectContextConfig, VibeConfig
|
||||
from vibe.core.tools.manager import ToolManager
|
||||
|
||||
|
||||
def _load_user_instructions() -> str:
|
||||
try:
|
||||
return INSTRUCTIONS_FILE.read_text("utf-8", errors="ignore")
|
||||
except (FileNotFoundError, OSError):
|
||||
return ""
|
||||
|
||||
|
||||
def _load_project_doc(workdir: Path, max_bytes: int) -> str:
|
||||
for name in PROJECT_DOC_FILENAMES:
|
||||
path = workdir / name
|
||||
try:
|
||||
return path.read_text("utf-8", errors="ignore")[:max_bytes]
|
||||
except (FileNotFoundError, OSError):
|
||||
continue
|
||||
return ""
|
||||
|
||||
|
||||
class ProjectContextProvider:
|
||||
def __init__(
|
||||
self, config: ProjectContextConfig, root_path: str | Path = "."
|
||||
) -> None:
|
||||
self.root_path = Path(root_path).resolve()
|
||||
self.config = config
|
||||
self.gitignore_patterns = self._load_gitignore_patterns()
|
||||
self._file_count = 0
|
||||
self._start_time = 0.0
|
||||
|
||||
def _load_gitignore_patterns(self) -> list[str]:
|
||||
gitignore_path = self.root_path / ".gitignore"
|
||||
patterns = []
|
||||
|
||||
if gitignore_path.exists():
|
||||
try:
|
||||
patterns.extend(
|
||||
line.strip()
|
||||
for line in gitignore_path.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip() and not line.startswith("#")
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not read .gitignore: {e}", file=sys.stderr)
|
||||
|
||||
default_patterns = [
|
||||
".git",
|
||||
".git/*",
|
||||
"*.pyc",
|
||||
"__pycache__",
|
||||
"node_modules",
|
||||
"node_modules/*",
|
||||
".env",
|
||||
".DS_Store",
|
||||
"*.log",
|
||||
".vscode/settings.json",
|
||||
".idea/*",
|
||||
"dist",
|
||||
"build",
|
||||
"target",
|
||||
".next",
|
||||
".nuxt",
|
||||
"coverage",
|
||||
".nyc_output",
|
||||
"*.egg-info",
|
||||
".pytest_cache",
|
||||
".tox",
|
||||
"vendor",
|
||||
"third_party",
|
||||
"deps",
|
||||
"*.min.js",
|
||||
"*.min.css",
|
||||
"*.bundle.js",
|
||||
"*.chunk.js",
|
||||
".cache",
|
||||
"tmp",
|
||||
"temp",
|
||||
"logs",
|
||||
]
|
||||
|
||||
return patterns + default_patterns
|
||||
|
||||
def _is_ignored(self, path: Path) -> bool:
|
||||
try:
|
||||
relative_path = path.relative_to(self.root_path)
|
||||
path_str = str(relative_path)
|
||||
|
||||
for pattern in self.gitignore_patterns:
|
||||
if pattern.endswith("/"):
|
||||
if path.is_dir() and fnmatch.fnmatch(f"{path_str}/", pattern):
|
||||
return True
|
||||
elif fnmatch.fnmatch(path_str, pattern):
|
||||
return True
|
||||
elif "*" in pattern or "?" in pattern:
|
||||
if fnmatch.fnmatch(path_str, pattern):
|
||||
return True
|
||||
|
||||
return False
|
||||
except (ValueError, OSError):
|
||||
return True
|
||||
|
||||
def _should_stop(self) -> bool:
|
||||
return (
|
||||
self._file_count >= self.config.max_files
|
||||
or (time.time() - self._start_time) > self.config.timeout_seconds
|
||||
)
|
||||
|
||||
def _build_tree_structure_iterative(self) -> Generator[str]:
|
||||
self._start_time = time.time()
|
||||
self._file_count = 0
|
||||
|
||||
yield from self._process_directory(self.root_path, "", 0, is_root=True)
|
||||
|
||||
def _process_directory(
|
||||
self, path: Path, prefix: str, depth: int, is_root: bool = False
|
||||
) -> Generator[str]:
|
||||
if depth > self.config.max_depth or self._should_stop():
|
||||
return
|
||||
|
||||
try:
|
||||
all_items = list(path.iterdir())
|
||||
items = [item for item in all_items if not self._is_ignored(item)]
|
||||
|
||||
items.sort(key=lambda p: (not p.is_dir(), p.name.lower()))
|
||||
|
||||
show_truncation = len(items) > self.config.max_dirs_per_level
|
||||
if show_truncation:
|
||||
items = items[: self.config.max_dirs_per_level]
|
||||
|
||||
for i, item in enumerate(items):
|
||||
if self._should_stop():
|
||||
break
|
||||
|
||||
is_last = i == len(items) - 1 and not show_truncation
|
||||
connector = "└── " if is_last else "├── "
|
||||
name = f"{item.name}{'/' if item.is_dir() else ''}"
|
||||
|
||||
yield f"{prefix}{connector}{name}"
|
||||
self._file_count += 1
|
||||
|
||||
if item.is_dir() and depth < self.config.max_depth:
|
||||
child_prefix = prefix + (" " if is_last else "│ ")
|
||||
yield from self._process_directory(item, child_prefix, depth + 1)
|
||||
|
||||
if show_truncation and not self._should_stop():
|
||||
remaining = len(all_items) - len(items)
|
||||
yield f"{prefix}└── ... ({remaining} more items)"
|
||||
|
||||
except (PermissionError, OSError):
|
||||
pass
|
||||
|
||||
def get_directory_structure(self) -> str:
|
||||
lines = []
|
||||
header = f"Directory structure of {self.root_path.name} (depth≤{self.config.max_depth}, max {self.config.max_files} items):\n"
|
||||
|
||||
try:
|
||||
for line in self._build_tree_structure_iterative():
|
||||
lines.append(line)
|
||||
|
||||
current_text = header + "\n".join(lines)
|
||||
if (
|
||||
len(current_text)
|
||||
> self.config.max_chars - self.config.truncation_buffer
|
||||
):
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
lines.append(f"Error building structure: {e}")
|
||||
|
||||
structure = header + "\n".join(lines)
|
||||
|
||||
if self._file_count >= self.config.max_files:
|
||||
structure += f"\n... (truncated at {self.config.max_files} files limit)"
|
||||
elif (time.time() - self._start_time) > self.config.timeout_seconds:
|
||||
structure += (
|
||||
f"\n... (truncated due to {self.config.timeout_seconds}s timeout)"
|
||||
)
|
||||
elif len(structure) > self.config.max_chars:
|
||||
structure += f"\n... (truncated at {self.config.max_chars} characters)"
|
||||
|
||||
return structure
|
||||
|
||||
def get_git_status(self) -> str:
|
||||
try:
|
||||
timeout = min(self.config.timeout_seconds, 10.0)
|
||||
num_commits = self.config.default_commit_count
|
||||
|
||||
current_branch = subprocess.run(
|
||||
["git", "branch", "--show-current"],
|
||||
capture_output=True,
|
||||
check=True,
|
||||
cwd=self.root_path,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
).stdout.strip()
|
||||
|
||||
main_branch = "main"
|
||||
try:
|
||||
branches_output = subprocess.run(
|
||||
["git", "branch", "-r"],
|
||||
capture_output=True,
|
||||
check=True,
|
||||
cwd=self.root_path,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
).stdout
|
||||
if "origin/master" in branches_output:
|
||||
main_branch = "master"
|
||||
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
status_output = subprocess.run(
|
||||
["git", "status", "--porcelain"],
|
||||
capture_output=True,
|
||||
check=True,
|
||||
cwd=self.root_path,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
).stdout.strip()
|
||||
|
||||
if status_output:
|
||||
status_lines = status_output.splitlines()
|
||||
MAX_GIT_STATUS_SIZE = 50
|
||||
if len(status_lines) > MAX_GIT_STATUS_SIZE:
|
||||
status = (
|
||||
f"({len(status_lines)} changes - use 'git status' for details)"
|
||||
)
|
||||
else:
|
||||
status = f"({len(status_lines)} changes)"
|
||||
else:
|
||||
status = "(clean)"
|
||||
|
||||
log_output = subprocess.run(
|
||||
["git", "log", "--oneline", f"-{num_commits}", "--decorate"],
|
||||
capture_output=True,
|
||||
check=True,
|
||||
cwd=self.root_path,
|
||||
stdin=subprocess.DEVNULL if is_windows() else None,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
).stdout.strip()
|
||||
|
||||
recent_commits = []
|
||||
for line in log_output.split("\n"):
|
||||
if not (line := line.strip()):
|
||||
continue
|
||||
|
||||
if " " in line:
|
||||
commit_hash, commit_msg = line.split(" ", 1)
|
||||
if (
|
||||
"(" in commit_msg
|
||||
and ")" in commit_msg
|
||||
and (paren_index := commit_msg.rfind("(")) > 0
|
||||
):
|
||||
commit_msg = commit_msg[:paren_index].strip()
|
||||
recent_commits.append(f"{commit_hash} {commit_msg}")
|
||||
else:
|
||||
recent_commits.append(line)
|
||||
|
||||
git_info_parts = [
|
||||
f"Current branch: {current_branch}",
|
||||
f"Main branch (you will usually use this for PRs): {main_branch}",
|
||||
f"Status: {status}",
|
||||
]
|
||||
|
||||
if recent_commits:
|
||||
git_info_parts.append("Recent commits:")
|
||||
git_info_parts.extend(recent_commits)
|
||||
|
||||
return "\n".join(git_info_parts)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return "Git operations timed out (large repository)"
|
||||
except subprocess.CalledProcessError:
|
||||
return "Not a git repository or git not available"
|
||||
except Exception as e:
|
||||
return f"Error getting git status: {e}"
|
||||
|
||||
def get_full_context(self) -> str:
|
||||
structure = self.get_directory_structure()
|
||||
git_status = self.get_git_status()
|
||||
|
||||
large_repo_warning = ""
|
||||
if len(structure) >= self.config.max_chars - self.config.truncation_buffer:
|
||||
large_repo_warning = (
|
||||
f" Large repository detected - showing summary view with depth limit {self.config.max_depth}. "
|
||||
f"Use the LS tool (passing a specific path), Bash tool, and other tools to explore nested directories in detail."
|
||||
)
|
||||
|
||||
template = UtilityPrompt.PROJECT_CONTEXT.read()
|
||||
return template.format(
|
||||
large_repo_warning=large_repo_warning,
|
||||
structure=structure,
|
||||
abs_path=self.root_path,
|
||||
git_status=git_status,
|
||||
)
|
||||
|
||||
|
||||
def _get_platform_name() -> str:
|
||||
platform_names = {
|
||||
"win32": "Windows",
|
||||
"darwin": "macOS",
|
||||
"linux": "Linux",
|
||||
"freebsd": "FreeBSD",
|
||||
"openbsd": "OpenBSD",
|
||||
"netbsd": "NetBSD",
|
||||
}
|
||||
return platform_names.get(sys.platform, "Unix-like")
|
||||
|
||||
|
||||
def _get_default_shell() -> str:
|
||||
"""Get the default shell used by asyncio.create_subprocess_shell.
|
||||
|
||||
On Unix, this is always 'sh'.
|
||||
On Windows, this is COMSPEC or cmd.exe.
|
||||
"""
|
||||
if is_windows():
|
||||
return os.environ.get("COMSPEC", "cmd.exe")
|
||||
return "sh"
|
||||
|
||||
|
||||
def _get_os_system_prompt() -> str:
|
||||
shell = _get_default_shell()
|
||||
platform_name = _get_platform_name()
|
||||
prompt = f"The operating system is {platform_name} with shell `{shell}`"
|
||||
|
||||
if is_windows():
|
||||
prompt += "\n" + _get_windows_system_prompt()
|
||||
return prompt
|
||||
|
||||
|
||||
def _get_windows_system_prompt() -> str:
|
||||
return (
|
||||
"### COMMAND COMPATIBILITY RULES (MUST FOLLOW):\n"
|
||||
"- DO NOT use Unix commands like `ls`, `grep`, `cat` - they won't work on Windows\n"
|
||||
"- Use: `dir` (Windows) for directory listings\n"
|
||||
"- Use: backslashes (\\\\) for paths\n"
|
||||
"- Check command availability with: `where command` (Windows)\n"
|
||||
"- Script shebang: Not applicable on Windows\n"
|
||||
"### ALWAYS verify commands work on the detected platform before suggesting them"
|
||||
)
|
||||
|
||||
|
||||
def get_universal_system_prompt(tool_manager: ToolManager, config: VibeConfig) -> str:
|
||||
sections = [config.system_prompt]
|
||||
|
||||
if config.include_model_info:
|
||||
sections.append(f"Your model name is: `{config.active_model}`")
|
||||
|
||||
if config.include_prompt_detail:
|
||||
sections.append(_get_os_system_prompt())
|
||||
tool_prompts = []
|
||||
active_tools = get_active_tool_classes(tool_manager, config)
|
||||
for tool_class in active_tools:
|
||||
if prompt := tool_class.get_tool_prompt():
|
||||
tool_prompts.append(prompt)
|
||||
if tool_prompts:
|
||||
sections.append("\n---\n".join(tool_prompts))
|
||||
|
||||
user_instructions = config.instructions.strip() or _load_user_instructions()
|
||||
if user_instructions.strip():
|
||||
sections.append(user_instructions)
|
||||
|
||||
if config.include_project_context:
|
||||
is_dangerous, reason = is_dangerous_directory()
|
||||
if is_dangerous:
|
||||
template = UtilityPrompt.DANGEROUS_DIRECTORY.read()
|
||||
context = template.format(
|
||||
reason=reason.lower(), abs_path=Path(".").resolve()
|
||||
)
|
||||
else:
|
||||
context = ProjectContextProvider(
|
||||
config=config.project_context, root_path=config.effective_workdir
|
||||
).get_full_context()
|
||||
|
||||
sections.append(context)
|
||||
|
||||
project_doc = _load_project_doc(
|
||||
config.effective_workdir, config.project_context.max_doc_bytes
|
||||
)
|
||||
if project_doc.strip():
|
||||
sections.append(project_doc)
|
||||
|
||||
return "\n\n".join(sections)
|
||||
281
vibe/core/tools/base.py
Normal file
281
vibe/core/tools/base.py
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import StrEnum, auto
|
||||
import functools
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sys
|
||||
from typing import Any, ClassVar, cast, get_args, get_type_hints
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
|
||||
|
||||
ARGS_COUNT = 4
|
||||
|
||||
|
||||
class ToolError(Exception):
|
||||
"""Raised when the tool encounters an unrecoverable problem."""
|
||||
|
||||
|
||||
class ToolInfo(BaseModel):
|
||||
"""Information about a tool.
|
||||
|
||||
Attributes:
|
||||
name: The name of the tool.
|
||||
description: A brief description of what the tool does.
|
||||
parameters: A dictionary of parameters required by the tool.
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
parameters: dict[str, Any]
|
||||
|
||||
|
||||
class ToolPermissionError(Exception):
|
||||
"""Raised when a tool permission is not allowed."""
|
||||
|
||||
|
||||
class ToolPermission(StrEnum):
|
||||
ALWAYS = auto()
|
||||
NEVER = auto()
|
||||
ASK = auto()
|
||||
|
||||
@classmethod
|
||||
def by_name(cls, name: str) -> ToolPermission:
|
||||
try:
|
||||
return ToolPermission(name.upper())
|
||||
except ValueError:
|
||||
raise ToolPermissionError(
|
||||
f"Invalid tool permission: {name}. Must be one of {list(cls)}"
|
||||
)
|
||||
|
||||
|
||||
class BaseToolConfig(BaseModel):
|
||||
"""Configuration for a tool.
|
||||
|
||||
Attributes:
|
||||
permission: The permission level required to use the tool.
|
||||
workdir: The working directory for the tool. If None, the current working directory is used.
|
||||
allowlist: Patterns that automatically allow tool execution.
|
||||
denylist: Patterns that automatically deny tool execution.
|
||||
"""
|
||||
|
||||
permission: ToolPermission = ToolPermission.ASK
|
||||
workdir: Path | None = Field(default=None, exclude=True)
|
||||
allowlist: list[str] = Field(default_factory=list)
|
||||
denylist: list[str] = Field(default_factory=list)
|
||||
|
||||
@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):
|
||||
return Path(v).expanduser().resolve()
|
||||
if isinstance(v, Path):
|
||||
return v.expanduser().resolve()
|
||||
return None
|
||||
|
||||
@property
|
||||
def effective_workdir(self) -> Path:
|
||||
return self.workdir if self.workdir is not None else Path.cwd()
|
||||
|
||||
|
||||
class BaseToolState(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
extra="forbid", validate_default=True, arbitrary_types_allowed=True
|
||||
)
|
||||
|
||||
|
||||
class BaseTool[
|
||||
ToolArgs: BaseModel,
|
||||
ToolResult: BaseModel,
|
||||
ToolConfig: BaseToolConfig,
|
||||
ToolState: BaseToolState,
|
||||
](ABC):
|
||||
description: ClassVar[str] = (
|
||||
"Base class for new tools. "
|
||||
"(Hey AI, if you're seeing this, someone skipped writing a description. "
|
||||
"Please gently meow at the developer to fix this.)"
|
||||
)
|
||||
|
||||
prompt_path: ClassVar[Path] | None = None
|
||||
|
||||
def __init__(self, config: ToolConfig, state: ToolState) -> None:
|
||||
self.config = config
|
||||
self.state = state
|
||||
|
||||
@abstractmethod
|
||||
async def run(self, args: ToolArgs) -> ToolResult:
|
||||
"""Invoke the tool with the given arguments. This method must be async."""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
@functools.cache
|
||||
def get_tool_prompt(cls) -> str | None:
|
||||
"""Loads and returns the content of the tool's .md prompt file, if it exists.
|
||||
|
||||
The prompt file is expected to be in a 'prompts' subdirectory relative to
|
||||
the tool's source file, with the same name but a .md extension
|
||||
(e.g., bash.py -> prompts/bash.md).
|
||||
"""
|
||||
try:
|
||||
class_file = inspect.getfile(cls)
|
||||
class_path = Path(class_file)
|
||||
prompt_dir = class_path.parent / "prompts"
|
||||
prompt_path = cls.prompt_path or prompt_dir / f"{class_path.stem}.md"
|
||||
|
||||
return prompt_path.read_text("utf-8")
|
||||
except (FileNotFoundError, TypeError, OSError):
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
async def invoke(self, **raw: Any) -> ToolResult:
|
||||
"""Validate arguments and run the tool.
|
||||
Pattern checking is now handled by Agent._should_execute_tool.
|
||||
"""
|
||||
try:
|
||||
args_model, _ = self._get_tool_args_results()
|
||||
args = args_model.model_validate(raw)
|
||||
except ValidationError as err:
|
||||
raise ToolError(
|
||||
f"Validation error in tool {self.get_name()}: {err}"
|
||||
) from err
|
||||
|
||||
return await self.run(args)
|
||||
|
||||
@classmethod
|
||||
def from_config(
|
||||
cls, config: ToolConfig
|
||||
) -> BaseTool[ToolArgs, ToolResult, ToolConfig, ToolState]:
|
||||
state_class = cls._get_tool_state_class()
|
||||
initial_state = state_class()
|
||||
return cls(config=config, state=initial_state)
|
||||
|
||||
@classmethod
|
||||
def _get_tool_config_class(cls) -> type[ToolConfig]:
|
||||
for base in getattr(cls, "__orig_bases__", ()):
|
||||
if getattr(base, "__origin__", None) is BaseTool:
|
||||
type_args = get_args(base)
|
||||
if len(type_args) == ARGS_COUNT:
|
||||
config_model = type_args[2]
|
||||
if issubclass(config_model, BaseToolConfig):
|
||||
return cast(type[ToolConfig], config_model)
|
||||
|
||||
for base_class in cls.__bases__:
|
||||
if base_class is object or base_class is ABC:
|
||||
continue
|
||||
try:
|
||||
return base_class._get_tool_config_class()
|
||||
except (TypeError, AttributeError):
|
||||
continue
|
||||
|
||||
raise TypeError(
|
||||
f"Could not determine ToolConfig for {cls.__name__}. "
|
||||
"Ensure it inherits from BaseTool with concrete type arguments."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_tool_state_class(cls) -> type[ToolState]:
|
||||
for base in getattr(cls, "__orig_bases__", ()):
|
||||
if getattr(base, "__origin__", None) is BaseTool:
|
||||
type_args = get_args(base)
|
||||
if len(type_args) == ARGS_COUNT:
|
||||
state_model = type_args[3]
|
||||
if issubclass(state_model, BaseToolState):
|
||||
return cast(type[ToolState], state_model)
|
||||
|
||||
for base_class in cls.__bases__:
|
||||
if base_class is object or base_class is ABC:
|
||||
continue
|
||||
try:
|
||||
return base_class._get_tool_state_class()
|
||||
except (TypeError, AttributeError):
|
||||
continue
|
||||
|
||||
raise TypeError(
|
||||
f"Could not determine ToolState for {cls.__name__}. "
|
||||
"Ensure it inherits from BaseTool with concrete type arguments."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_tool_args_results(cls) -> tuple[type[ToolArgs], type[ToolResult]]:
|
||||
"""Extract <ToolArgs, ToolResult> from the annotated signature of `run`.
|
||||
Works even when `from __future__ import annotations` is in effect.
|
||||
"""
|
||||
run_fn = cls.run.__func__ if isinstance(cls.run, classmethod) else cls.run
|
||||
|
||||
type_hints = get_type_hints(
|
||||
run_fn,
|
||||
globalns=vars(sys.modules[cls.__module__]),
|
||||
localns={cls.__name__: cls},
|
||||
)
|
||||
|
||||
try:
|
||||
args_model = type_hints["args"]
|
||||
result_model = type_hints["return"]
|
||||
except KeyError as e:
|
||||
raise TypeError(
|
||||
f"{cls.__name__}.run must be annotated as "
|
||||
"`async def run(self, args: ToolArgs) -> ToolResult`"
|
||||
) from e
|
||||
|
||||
if not (
|
||||
issubclass(args_model, BaseModel) and issubclass(result_model, BaseModel)
|
||||
):
|
||||
raise TypeError(
|
||||
f"{cls.__name__}.run annotations must be Pydantic models; "
|
||||
f"got {args_model!r}, {result_model!r}"
|
||||
)
|
||||
|
||||
return cast(type[ToolArgs], args_model), cast(type[ToolResult], result_model)
|
||||
|
||||
@classmethod
|
||||
def get_parameters(cls) -> dict[str, Any]:
|
||||
"""Return a cleaned-up JSON-schema dict describing the arguments model
|
||||
with which this concrete tool was parametrised.
|
||||
"""
|
||||
args_model, _ = cls._get_tool_args_results()
|
||||
schema = args_model.model_json_schema()
|
||||
schema.pop("title", None)
|
||||
schema.pop("description", None)
|
||||
|
||||
if "properties" in schema:
|
||||
for prop_details in schema["properties"].values():
|
||||
prop_details.pop("title", None)
|
||||
|
||||
if "$defs" in schema:
|
||||
for def_details in schema["$defs"].values():
|
||||
def_details.pop("title", None)
|
||||
if "properties" in def_details:
|
||||
for prop_details in def_details["properties"].values():
|
||||
prop_details.pop("title", None)
|
||||
|
||||
return schema
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
name = cls.__name__
|
||||
snake_case = re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
|
||||
return snake_case
|
||||
|
||||
@classmethod
|
||||
def create_config_with_permission(
|
||||
cls, permission: ToolPermission
|
||||
) -> BaseToolConfig:
|
||||
config_class = cls._get_tool_config_class()
|
||||
return config_class(permission=permission)
|
||||
|
||||
def check_allowlist_denylist(self, args: ToolArgs) -> ToolPermission | None:
|
||||
"""Check if args match allowlist/denylist patterns.
|
||||
|
||||
Returns:
|
||||
ToolPermission.ALWAYS if allowlisted
|
||||
ToolPermission.NEVER if denylisted
|
||||
None if no match (proceed with normal permission check)
|
||||
|
||||
Base implementation returns None. Override in subclasses for specific logic.
|
||||
"""
|
||||
return None
|
||||
285
vibe/core/tools/builtins/bash.py
Normal file
285
vibe/core/tools/builtins/bash.py
Normal file
|
|
@ -0,0 +1,285 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import sys
|
||||
from typing import ClassVar, Literal, final
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
BaseToolConfig,
|
||||
BaseToolState,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
)
|
||||
from vibe.core.utils import is_windows
|
||||
|
||||
|
||||
def _get_subprocess_encoding() -> str:
|
||||
if sys.platform == "win32":
|
||||
# Windows console uses OEM code page (e.g., cp850, cp1252)
|
||||
import ctypes
|
||||
|
||||
return f"cp{ctypes.windll.kernel32.GetOEMCP()}"
|
||||
return "utf-8"
|
||||
|
||||
|
||||
def _get_base_env() -> dict[str, str]:
|
||||
base_env = {
|
||||
**os.environ,
|
||||
"CI": "true",
|
||||
"NONINTERACTIVE": "1",
|
||||
"NO_TTY": "1",
|
||||
"NO_COLOR": "1",
|
||||
}
|
||||
|
||||
if is_windows():
|
||||
base_env["GIT_PAGER"] = "more"
|
||||
base_env["PAGER"] = "more"
|
||||
else:
|
||||
base_env["TERM"] = "dumb"
|
||||
base_env["DEBIAN_FRONTEND"] = "noninteractive"
|
||||
base_env["GIT_PAGER"] = "cat"
|
||||
base_env["PAGER"] = "cat"
|
||||
base_env["LESS"] = "-FX"
|
||||
base_env["LC_ALL"] = "en_US.UTF-8"
|
||||
|
||||
return base_env
|
||||
|
||||
|
||||
async def _kill_process_tree(proc: asyncio.subprocess.Process) -> None:
|
||||
if proc.returncode is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
try:
|
||||
subprocess_proc = await asyncio.create_subprocess_exec(
|
||||
"taskkill",
|
||||
"/F",
|
||||
"/T",
|
||||
"/PID",
|
||||
str(proc.pid),
|
||||
stdout=asyncio.subprocess.DEVNULL,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
)
|
||||
await subprocess_proc.wait()
|
||||
except (FileNotFoundError, OSError):
|
||||
proc.terminate()
|
||||
else:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
|
||||
await proc.wait()
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
pass
|
||||
|
||||
|
||||
def _get_default_allowlist() -> list[str]:
|
||||
common = ["echo", "find", "git diff", "git log", "git status", "tree", "whoami"]
|
||||
|
||||
if is_windows():
|
||||
return common + ["dir", "findstr", "more", "type", "ver", "where"]
|
||||
else:
|
||||
return common + [
|
||||
"cat",
|
||||
"file",
|
||||
"head",
|
||||
"ls",
|
||||
"pwd",
|
||||
"stat",
|
||||
"tail",
|
||||
"uname",
|
||||
"wc",
|
||||
"which",
|
||||
]
|
||||
|
||||
|
||||
def _get_default_denylist() -> list[str]:
|
||||
common = ["gdb", "pdb", "passwd"]
|
||||
|
||||
if is_windows():
|
||||
return common + ["cmd /k", "powershell -NoExit", "pwsh -NoExit", "notepad"]
|
||||
else:
|
||||
return common + [
|
||||
"nano",
|
||||
"vim",
|
||||
"vi",
|
||||
"emacs",
|
||||
"bash -i",
|
||||
"sh -i",
|
||||
"zsh -i",
|
||||
"fish -i",
|
||||
"dash -i",
|
||||
"screen",
|
||||
"tmux",
|
||||
]
|
||||
|
||||
|
||||
def _get_default_denylist_standalone() -> list[str]:
|
||||
common = ["python", "python3", "ipython"]
|
||||
|
||||
if is_windows():
|
||||
return common + ["cmd", "powershell", "pwsh", "notepad"]
|
||||
else:
|
||||
return common + ["bash", "sh", "nohup", "vi", "vim", "emacs", "nano", "su"]
|
||||
|
||||
|
||||
class BashToolConfig(BaseToolConfig):
|
||||
permission: ToolPermission = ToolPermission.ASK
|
||||
max_output_bytes: int = Field(
|
||||
default=16_000, description="Maximum bytes to capture from stdout and stderr."
|
||||
)
|
||||
default_timeout: int = Field(
|
||||
default=30, description="Default timeout for commands in seconds."
|
||||
)
|
||||
allowlist: list[str] = Field(
|
||||
default_factory=_get_default_allowlist,
|
||||
description="Command prefixes that are automatically allowed",
|
||||
)
|
||||
denylist: list[str] = Field(
|
||||
default_factory=_get_default_denylist,
|
||||
description="Command prefixes that are automatically denied",
|
||||
)
|
||||
denylist_standalone: list[str] = Field(
|
||||
default_factory=_get_default_denylist_standalone,
|
||||
description="Commands that are denied only when run without arguments",
|
||||
)
|
||||
|
||||
|
||||
class BashArgs(BaseModel):
|
||||
command: str
|
||||
timeout: int | None = Field(
|
||||
default=None, description="Override the default command timeout."
|
||||
)
|
||||
|
||||
|
||||
class BashResult(BaseModel):
|
||||
stdout: str
|
||||
stderr: str
|
||||
returncode: int
|
||||
|
||||
|
||||
class Bash(BaseTool[BashArgs, BashResult, BashToolConfig, BaseToolState]):
|
||||
description: ClassVar[str] = "Run a one-off bash command and capture its output."
|
||||
|
||||
def check_allowlist_denylist(self, args: BashArgs) -> ToolPermission | None:
|
||||
command_parts = re.split(r"(?:&&|\|\||;|\|)", args.command)
|
||||
command_parts = [part.strip() for part in command_parts if part.strip()]
|
||||
|
||||
if not command_parts:
|
||||
return None
|
||||
|
||||
def is_denylisted(command: str) -> bool:
|
||||
return any(command.startswith(pattern) for pattern in self.config.denylist)
|
||||
|
||||
def is_standalone_denylisted(command: str) -> bool:
|
||||
parts = command.split()
|
||||
if not parts:
|
||||
return False
|
||||
|
||||
base_command = parts[0]
|
||||
has_args = len(parts) > 1
|
||||
|
||||
if not has_args:
|
||||
command_name = os.path.basename(base_command)
|
||||
if command_name in self.config.denylist_standalone:
|
||||
return True
|
||||
if base_command in self.config.denylist_standalone:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def is_allowlisted(command: str) -> bool:
|
||||
return any(command.startswith(pattern) for pattern in self.config.allowlist)
|
||||
|
||||
for part in command_parts:
|
||||
if is_denylisted(part):
|
||||
return ToolPermission.NEVER
|
||||
if is_standalone_denylisted(part):
|
||||
return ToolPermission.NEVER
|
||||
|
||||
if all(is_allowlisted(part) for part in command_parts):
|
||||
return ToolPermission.ALWAYS
|
||||
|
||||
return None
|
||||
|
||||
@final
|
||||
def _build_timeout_error(self, command: str, timeout: int) -> ToolError:
|
||||
return ToolError(f"Command timed out after {timeout}s: {command!r}")
|
||||
|
||||
@final
|
||||
def _build_result(
|
||||
self, *, command: str, stdout: str, stderr: str, returncode: int
|
||||
) -> BashResult:
|
||||
if returncode != 0:
|
||||
error_msg = f"Command failed: {command!r}\n"
|
||||
error_msg += f"Return code: {returncode}"
|
||||
if stderr:
|
||||
error_msg += f"\nStderr: {stderr}"
|
||||
if stdout:
|
||||
error_msg += f"\nStdout: {stdout}"
|
||||
raise ToolError(error_msg.strip())
|
||||
|
||||
return BashResult(stdout=stdout, stderr=stderr, returncode=returncode)
|
||||
|
||||
async def run(self, args: BashArgs) -> BashResult:
|
||||
timeout = args.timeout or self.config.default_timeout
|
||||
max_bytes = self.config.max_output_bytes
|
||||
|
||||
proc = None
|
||||
try:
|
||||
# start_new_session is Unix-only, on Windows it's ignored
|
||||
kwargs: dict[Literal["start_new_session"], bool] = (
|
||||
{} if is_windows() else {"start_new_session": True}
|
||||
)
|
||||
|
||||
proc = await asyncio.create_subprocess_shell(
|
||||
args.command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
cwd=self.config.effective_workdir,
|
||||
env=_get_base_env(),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
proc.communicate(), timeout=timeout
|
||||
)
|
||||
except TimeoutError:
|
||||
await _kill_process_tree(proc)
|
||||
raise self._build_timeout_error(args.command, timeout)
|
||||
|
||||
encoding = _get_subprocess_encoding()
|
||||
stdout = (
|
||||
stdout_bytes.decode(encoding, errors="replace")[:max_bytes]
|
||||
if stdout_bytes
|
||||
else ""
|
||||
)
|
||||
stderr = (
|
||||
stderr_bytes.decode(encoding, errors="replace")[:max_bytes]
|
||||
if stderr_bytes
|
||||
else ""
|
||||
)
|
||||
|
||||
returncode = proc.returncode or 0
|
||||
|
||||
return self._build_result(
|
||||
command=args.command,
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
returncode=returncode,
|
||||
)
|
||||
|
||||
except (ToolError, asyncio.CancelledError):
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ToolError(f"Error running command {args.command!r}: {exc}") from exc
|
||||
finally:
|
||||
if proc is not None:
|
||||
await _kill_process_tree(proc)
|
||||
326
vibe/core/tools/builtins/grep.py
Normal file
326
vibe/core/tools/builtins/grep.py
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from enum import StrEnum, auto
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
BaseToolConfig,
|
||||
BaseToolState,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
)
|
||||
from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class GrepBackend(StrEnum):
|
||||
RIPGREP = auto()
|
||||
GNU_GREP = auto()
|
||||
|
||||
|
||||
class GrepToolConfig(BaseToolConfig):
|
||||
permission: ToolPermission = ToolPermission.ALWAYS
|
||||
|
||||
max_output_bytes: int = Field(
|
||||
default=64_000, description="Hard cap for the total size of matched lines."
|
||||
)
|
||||
default_max_matches: int = Field(
|
||||
default=100, description="Default maximum number of matches to return."
|
||||
)
|
||||
default_timeout: int = Field(
|
||||
default=60, description="Default timeout for the search command in seconds."
|
||||
)
|
||||
exclude_patterns: list[str] = Field(
|
||||
default=[
|
||||
".venv/",
|
||||
"venv/",
|
||||
".env/",
|
||||
"env/",
|
||||
"node_modules/",
|
||||
".git/",
|
||||
"__pycache__/",
|
||||
".pytest_cache/",
|
||||
".mypy_cache/",
|
||||
".tox/",
|
||||
".nox/",
|
||||
".coverage/",
|
||||
"htmlcov/",
|
||||
"dist/",
|
||||
"build/",
|
||||
".idea/",
|
||||
".vscode/",
|
||||
"*.egg-info",
|
||||
"*.pyc",
|
||||
"*.pyo",
|
||||
"*.pyd",
|
||||
".DS_Store",
|
||||
"Thumbs.db",
|
||||
],
|
||||
description="List of glob patterns to exclude from search (dirs should end with /).",
|
||||
)
|
||||
codeignore_file: str = Field(
|
||||
default=".vibeignore",
|
||||
description="Name of the file to read for additional exclusion patterns.",
|
||||
)
|
||||
|
||||
|
||||
class GrepState(BaseToolState):
|
||||
search_history: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class GrepArgs(BaseModel):
|
||||
pattern: str
|
||||
path: str = "."
|
||||
max_matches: int | None = Field(
|
||||
default=None, description="Override the default maximum number of matches."
|
||||
)
|
||||
use_default_ignore: bool = Field(
|
||||
default=True, description="Whether to respect .gitignore and .ignore files."
|
||||
)
|
||||
|
||||
|
||||
class GrepResult(BaseModel):
|
||||
matches: str
|
||||
match_count: int
|
||||
was_truncated: bool = Field(
|
||||
description="True if output was cut short by max_matches or max_output_bytes."
|
||||
)
|
||||
|
||||
|
||||
class Grep(
|
||||
BaseTool[GrepArgs, GrepResult, GrepToolConfig, GrepState],
|
||||
ToolUIData[GrepArgs, GrepResult],
|
||||
):
|
||||
description: ClassVar[str] = (
|
||||
"Recursively search files for a regex pattern using ripgrep (rg) or grep. "
|
||||
"Respects .gitignore and .codeignore files by default when using ripgrep."
|
||||
)
|
||||
|
||||
def _detect_backend(self) -> GrepBackend:
|
||||
if shutil.which("rg"):
|
||||
return GrepBackend.RIPGREP
|
||||
if shutil.which("grep"):
|
||||
return GrepBackend.GNU_GREP
|
||||
raise ToolError(
|
||||
"Neither ripgrep (rg) nor grep is installed. "
|
||||
"Please install ripgrep: https://github.com/BurntSushi/ripgrep#installation"
|
||||
)
|
||||
|
||||
async def run(self, args: GrepArgs) -> GrepResult:
|
||||
backend = self._detect_backend()
|
||||
self._validate_args(args)
|
||||
self.state.search_history.append(args.pattern)
|
||||
|
||||
exclude_patterns = self._collect_exclude_patterns()
|
||||
cmd = self._build_command(args, exclude_patterns, backend)
|
||||
stdout = await self._execute_search(cmd)
|
||||
|
||||
return self._parse_output(
|
||||
stdout, args.max_matches or self.config.default_max_matches
|
||||
)
|
||||
|
||||
def _validate_args(self, args: GrepArgs) -> None:
|
||||
if not args.pattern.strip():
|
||||
raise ToolError("Empty search pattern provided.")
|
||||
|
||||
path_obj = Path(args.path).expanduser()
|
||||
if not path_obj.is_absolute():
|
||||
path_obj = self.config.effective_workdir / path_obj
|
||||
|
||||
if not path_obj.exists():
|
||||
raise ToolError(f"Path does not exist: {args.path}")
|
||||
|
||||
def _collect_exclude_patterns(self) -> list[str]:
|
||||
patterns = list(self.config.exclude_patterns)
|
||||
|
||||
codeignore_path = self.config.effective_workdir / self.config.codeignore_file
|
||||
if codeignore_path.is_file():
|
||||
patterns.extend(self._load_codeignore_patterns(codeignore_path))
|
||||
|
||||
return patterns
|
||||
|
||||
def _load_codeignore_patterns(self, codeignore_path: Path) -> list[str]:
|
||||
patterns = []
|
||||
try:
|
||||
content = codeignore_path.read_text("utf-8")
|
||||
for line in content.splitlines():
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#"):
|
||||
patterns.append(line)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return patterns
|
||||
|
||||
def _build_command(
|
||||
self, args: GrepArgs, exclude_patterns: list[str], backend: GrepBackend
|
||||
) -> list[str]:
|
||||
if backend == GrepBackend.RIPGREP:
|
||||
return self._build_ripgrep_command(args, exclude_patterns)
|
||||
return self._build_gnu_grep_command(args, exclude_patterns)
|
||||
|
||||
def _build_ripgrep_command(
|
||||
self, args: GrepArgs, exclude_patterns: list[str]
|
||||
) -> list[str]:
|
||||
max_matches = args.max_matches or self.config.default_max_matches
|
||||
|
||||
cmd = [
|
||||
"rg",
|
||||
"--line-number",
|
||||
"--no-heading",
|
||||
"--smart-case",
|
||||
"--no-binary",
|
||||
# Request one extra to detect truncation
|
||||
"--max-count",
|
||||
str(max_matches + 1),
|
||||
]
|
||||
|
||||
if not args.use_default_ignore:
|
||||
cmd.append("--no-ignore")
|
||||
|
||||
for pattern in exclude_patterns:
|
||||
cmd.extend(["--glob", f"!{pattern}"])
|
||||
|
||||
cmd.extend(["-e", args.pattern, args.path])
|
||||
|
||||
return cmd
|
||||
|
||||
def _build_gnu_grep_command(
|
||||
self, args: GrepArgs, exclude_patterns: list[str]
|
||||
) -> list[str]:
|
||||
max_matches = args.max_matches or self.config.default_max_matches
|
||||
|
||||
cmd = ["grep", "-r", "-n", "-I", "-E", f"--max-count={max_matches + 1}"]
|
||||
|
||||
if args.pattern.islower():
|
||||
cmd.append("-i")
|
||||
|
||||
for pattern in exclude_patterns:
|
||||
if pattern.endswith("/"):
|
||||
dir_pattern = pattern.rstrip("/")
|
||||
cmd.append(f"--exclude-dir={dir_pattern}")
|
||||
else:
|
||||
cmd.append(f"--exclude={pattern}")
|
||||
|
||||
cmd.extend(["-e", args.pattern, args.path])
|
||||
|
||||
return cmd
|
||||
|
||||
async def _execute_search(self, cmd: list[str]) -> str:
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=str(self.config.effective_workdir),
|
||||
)
|
||||
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
proc.communicate(), timeout=self.config.default_timeout
|
||||
)
|
||||
except TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
raise ToolError(
|
||||
f"Search timed out after {self.config.default_timeout}s"
|
||||
)
|
||||
|
||||
stdout = (
|
||||
stdout_bytes.decode("utf-8", errors="ignore") if stdout_bytes else ""
|
||||
)
|
||||
stderr = (
|
||||
stderr_bytes.decode("utf-8", errors="ignore") if stderr_bytes else ""
|
||||
)
|
||||
|
||||
if proc.returncode not in {0, 1}:
|
||||
error_msg = stderr or f"Process exited with code {proc.returncode}"
|
||||
raise ToolError(f"grep error: {error_msg}")
|
||||
|
||||
return stdout
|
||||
|
||||
except ToolError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ToolError(f"Error running grep: {exc}") from exc
|
||||
|
||||
def _parse_output(self, stdout: str, max_matches: int) -> GrepResult:
|
||||
output_lines = stdout.splitlines() if stdout else []
|
||||
|
||||
truncated_lines = output_lines[:max_matches]
|
||||
truncated_output = "\n".join(truncated_lines)
|
||||
|
||||
was_truncated = (
|
||||
len(output_lines) > max_matches
|
||||
or len(truncated_output) > self.config.max_output_bytes
|
||||
)
|
||||
|
||||
final_output = truncated_output[: self.config.max_output_bytes]
|
||||
|
||||
return GrepResult(
|
||||
matches=final_output,
|
||||
match_count=len(truncated_lines),
|
||||
was_truncated=was_truncated,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
|
||||
if not isinstance(event.args, GrepArgs):
|
||||
return ToolCallDisplay(summary="grep")
|
||||
|
||||
summary = f"grep: '{event.args.pattern}'"
|
||||
if event.args.path != ".":
|
||||
summary += f" in {event.args.path}"
|
||||
if event.args.max_matches:
|
||||
summary += f" (max {event.args.max_matches} matches)"
|
||||
if not event.args.use_default_ignore:
|
||||
summary += " [no-ignore]"
|
||||
|
||||
return ToolCallDisplay(
|
||||
summary=summary,
|
||||
details={
|
||||
"pattern": event.args.pattern,
|
||||
"path": event.args.path,
|
||||
"max_matches": event.args.max_matches,
|
||||
"use_default_ignore": event.args.use_default_ignore,
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
|
||||
if not isinstance(event.result, GrepResult):
|
||||
return ToolResultDisplay(
|
||||
success=False, message=event.error or event.skip_reason or "No result"
|
||||
)
|
||||
|
||||
message = f"Found {event.result.match_count} matches"
|
||||
if event.result.was_truncated:
|
||||
message += " (truncated)"
|
||||
|
||||
warnings = []
|
||||
if event.result.was_truncated:
|
||||
warnings.append("Output was truncated due to size/match limits")
|
||||
|
||||
return ToolResultDisplay(
|
||||
success=True,
|
||||
message=message,
|
||||
warnings=warnings,
|
||||
details={
|
||||
"match_count": event.result.match_count,
|
||||
"was_truncated": event.result.was_truncated,
|
||||
"matches": event.result.matches,
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_status_text(cls) -> str:
|
||||
"""Return status message for spinner."""
|
||||
return "Searching files"
|
||||
0
vibe/core/tools/builtins/prompts/__init__.py
Normal file
0
vibe/core/tools/builtins/prompts/__init__.py
Normal file
68
vibe/core/tools/builtins/prompts/bash.md
Normal file
68
vibe/core/tools/builtins/prompts/bash.md
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
Use the `bash` tool to run one-off shell commands.
|
||||
|
||||
**Key characteristics:**
|
||||
- **Stateless**: Each command runs independently in a fresh environment
|
||||
|
||||
**IMPORTANT: Use dedicated tools if available instead of these bash commands:**
|
||||
|
||||
**File Operations - DO NOT USE:**
|
||||
- `cat filename` → Use `read_file(path="filename")`
|
||||
- `head -n 20 filename` → Use `read_file(path="filename", limit=20)`
|
||||
- `tail -n 20 filename` → Read with offset: `read_file(path="filename", offset=<line_number>, limit=20)`
|
||||
- `sed -n '100,200p' filename` → Use `read_file(path="filename", offset=99, limit=101)`
|
||||
- `less`, `more`, `vim`, `nano` → Use `read_file` with offset/limit for navigation
|
||||
- `echo "content" > file` → Use `write_file(path="file", content="content")`
|
||||
- `echo "content" >> file` → Read first, then `write_file` with overwrite=true
|
||||
|
||||
**Search Operations - DO NOT USE:**
|
||||
- `grep -r "pattern" .` → Use `grep(pattern="pattern", path=".")`
|
||||
- `find . -name "*.py"` → Use `bash("ls -la")` for current dir or `grep` with appropriate pattern
|
||||
- `ag`, `ack`, `rg` commands → Use the `grep` tool
|
||||
- `locate` → Use `grep` tool
|
||||
|
||||
**File Modification - DO NOT USE:**
|
||||
- `sed -i 's/old/new/g' file` → Use `search_replace` tool
|
||||
- `awk` for file editing → Use `search_replace` tool
|
||||
- Any in-place file editing → Use `search_replace` tool
|
||||
|
||||
**APPROPRIATE bash uses:**
|
||||
- System information: `pwd`, `whoami`, `date`, `uname -a`
|
||||
- Directory listings: `ls -la`, `tree` (if available)
|
||||
- Git operations: `git status`, `git log --oneline -10`, `git diff`
|
||||
- Process info: `ps aux | grep process`, `top -n 1`
|
||||
- Network checks: `ping -c 1 google.com`, `curl -I https://example.com`
|
||||
- Package management: `pip list`, `npm list`
|
||||
- Environment checks: `env | grep VAR`, `which python`
|
||||
- File metadata: `stat filename`, `file filename`, `wc -l filename`
|
||||
|
||||
**Example: Reading a large file efficiently**
|
||||
|
||||
WRONG:
|
||||
```bash
|
||||
bash("cat large_file.txt") # May hit size limits
|
||||
bash("head -1000 large_file.txt") # Inefficient
|
||||
```
|
||||
|
||||
RIGHT:
|
||||
```python
|
||||
# First chunk
|
||||
read_file(path="large_file.txt", limit=1000)
|
||||
# If was_truncated=true, read next chunk
|
||||
read_file(path="large_file.txt", offset=1000, limit=1000)
|
||||
```
|
||||
|
||||
**Example: Searching for patterns**
|
||||
|
||||
WRONG:
|
||||
```bash
|
||||
bash("grep -r 'TODO' src/") # Don't use bash for grep
|
||||
bash("find . -type f -name '*.py' | xargs grep 'import'") # Too complex
|
||||
```
|
||||
|
||||
RIGHT:
|
||||
```python
|
||||
grep(pattern="TODO", path="src/")
|
||||
grep(pattern="import", path=".")
|
||||
```
|
||||
|
||||
**Remember:** Bash is best for quick system checks and git operations. For file operations, searching, and editing, always use the dedicated tools when they are available.
|
||||
4
vibe/core/tools/builtins/prompts/grep.md
Normal file
4
vibe/core/tools/builtins/prompts/grep.md
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
Use `grep` to recursively search for a regular expression pattern in files.
|
||||
|
||||
- It's very fast and automatically ignores files that you should not read like .pyc files, .venv directories, etc.
|
||||
- Use this to find where functions are defined, how variables are used, or to locate specific error messages.
|
||||
13
vibe/core/tools/builtins/prompts/read_file.md
Normal file
13
vibe/core/tools/builtins/prompts/read_file.md
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
Use `read_file` to read the content of a file. It's designed to handle large files safely.
|
||||
|
||||
- By default, it reads from the beginning of the file.
|
||||
- Use `offset` (line number) and `limit` (number of lines) to read specific parts or chunks of a file. This is efficient for exploring large files.
|
||||
- The result includes `was_truncated: true` if the file content was cut short due to size limits.
|
||||
|
||||
**Strategy for large files:**
|
||||
|
||||
1. Call `read_file` with a `limit` (e.g., 1000 lines) to get the start of the file.
|
||||
2. If `was_truncated` is true, you know the file is large.
|
||||
3. To read the next chunk, call `read_file` again with an `offset`. For example, `offset=1000, limit=1000`.
|
||||
|
||||
This is more efficient than using `bash` with `cat` or `wc`.
|
||||
43
vibe/core/tools/builtins/prompts/search_replace.md
Normal file
43
vibe/core/tools/builtins/prompts/search_replace.md
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
Use `search_replace` to make targeted changes to files using SEARCH/REPLACE blocks. This tool finds exact text matches and replaces them.
|
||||
|
||||
Arguments:
|
||||
- `file_path`: The path to the file to modify
|
||||
- `content`: The SEARCH/REPLACE blocks defining the changes
|
||||
|
||||
The content format is:
|
||||
|
||||
```
|
||||
<<<<<<< SEARCH
|
||||
[exact text to find in the file]
|
||||
=======
|
||||
[exact text to replace it with]
|
||||
>>>>>>> REPLACE
|
||||
```
|
||||
|
||||
You can include multiple SEARCH/REPLACE blocks to make multiple changes to the same file:
|
||||
|
||||
```
|
||||
<<<<<<< SEARCH
|
||||
def old_function():
|
||||
return "old value"
|
||||
=======
|
||||
def new_function():
|
||||
return "new value"
|
||||
>>>>>>> REPLACE
|
||||
|
||||
<<<<<<< SEARCH
|
||||
import os
|
||||
=======
|
||||
import os
|
||||
import sys
|
||||
>>>>>>> REPLACE
|
||||
```
|
||||
|
||||
IMPORTANT:
|
||||
|
||||
- The SEARCH text must match EXACTLY (including whitespace, indentation, and line endings)
|
||||
- The SEARCH text must appear exactly once in the file - if it appears multiple times, the tool will error
|
||||
- Use at least 5 equals signs (=====) between SEARCH and REPLACE sections
|
||||
- The tool will provide detailed error messages showing context if search text is not found
|
||||
- Each search/replace block is applied in order, so later blocks see the results of earlier ones
|
||||
- Be careful with escape sequences in string literals - use \n not \\n for newlines in code
|
||||
199
vibe/core/tools/builtins/prompts/todo.md
Normal file
199
vibe/core/tools/builtins/prompts/todo.md
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
Use the `todo` tool to manage a simple task list. This tool helps you track tasks and their progress.
|
||||
|
||||
## How it works
|
||||
|
||||
- **Reading:** Use `action: "read"` to view the current todo list
|
||||
- **Writing:** Use `action: "write"` with the complete `todos` list to update. You must provide the ENTIRE list - this replaces everything.
|
||||
|
||||
## Todo Structure
|
||||
Each todo item has:
|
||||
- `id`: A unique string identifier (e.g., "1", "2", "task-a")
|
||||
- `content`: The task description
|
||||
- `status`: One of: "pending", "in_progress", "completed", "cancelled"
|
||||
- `priority`: One of: "high", "medium", "low"
|
||||
|
||||
## When to Use This Tool
|
||||
|
||||
**Use proactively for:**
|
||||
- Complex multi-step tasks (3+ distinct steps)
|
||||
- Non-trivial tasks requiring careful planning
|
||||
- Multiple tasks provided by the user (numbered or comma-separated)
|
||||
- Tracking progress on ongoing work
|
||||
- After receiving new instructions - immediately capture requirements
|
||||
- When starting work - mark task as in_progress BEFORE beginning
|
||||
- After completing work - mark as completed and add any follow-up tasks discovered
|
||||
|
||||
**Skip this tool for:**
|
||||
- Single, straightforward tasks
|
||||
- Trivial operations (< 3 simple steps)
|
||||
- Purely conversational or informational requests
|
||||
- Tasks that provide no organizational benefit
|
||||
|
||||
## Task Management Best Practices
|
||||
|
||||
1. **Status Management:**
|
||||
- Only ONE task should be `in_progress` at a time
|
||||
- Mark tasks `in_progress` BEFORE starting work on them
|
||||
- Mark tasks `completed` IMMEDIATELY after finishing
|
||||
- Keep tasks `in_progress` if blocked or encountering errors
|
||||
|
||||
2. **Task Completion Rules:**
|
||||
- ONLY mark as `completed` when FULLY accomplished
|
||||
- Never mark complete if tests are failing, implementation is partial, or errors are unresolved
|
||||
- When blocked, create a new task describing what needs resolution
|
||||
|
||||
3. **Task Organization:**
|
||||
- Create specific, actionable items
|
||||
- Break complex tasks into manageable steps
|
||||
- Use clear, descriptive task names
|
||||
- Remove irrelevant tasks entirely (don't just mark cancelled)
|
||||
|
||||
## Examples
|
||||
|
||||
**Example 1: Reading todos**
|
||||
```json
|
||||
{
|
||||
"action": "read"
|
||||
}
|
||||
```
|
||||
|
||||
**Example 2: Initial task creation (user requests multiple features)**
|
||||
```json
|
||||
{
|
||||
"action": "write",
|
||||
"todos": [
|
||||
{
|
||||
"id": "1",
|
||||
"content": "Add dark mode toggle to settings",
|
||||
"status": "pending",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "2",
|
||||
"content": "Implement theme context/state management",
|
||||
"status": "pending",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "3",
|
||||
"content": "Update components for theme switching",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
},
|
||||
{
|
||||
"id": "4",
|
||||
"content": "Run tests and verify build",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Example 3: Starting work (marking one task in_progress)**
|
||||
```json
|
||||
{
|
||||
"action": "write",
|
||||
"todos": [
|
||||
{
|
||||
"id": "1",
|
||||
"content": "Add dark mode toggle to settings",
|
||||
"status": "in_progress",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "2",
|
||||
"content": "Implement theme context/state management",
|
||||
"status": "pending",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "3",
|
||||
"content": "Update components for theme switching",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
},
|
||||
{
|
||||
"id": "4",
|
||||
"content": "Run tests and verify build",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Example 4: Completing task and adding discovered subtask**
|
||||
```json
|
||||
{
|
||||
"action": "write",
|
||||
"todos": [
|
||||
{
|
||||
"id": "1",
|
||||
"content": "Add dark mode toggle to settings",
|
||||
"status": "completed",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "2",
|
||||
"content": "Implement theme context/state management",
|
||||
"status": "in_progress",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "3",
|
||||
"content": "Update components for theme switching",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
},
|
||||
{
|
||||
"id": "4",
|
||||
"content": "Fix TypeScript errors in theme types",
|
||||
"status": "pending",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "5",
|
||||
"content": "Run tests and verify build",
|
||||
"status": "pending",
|
||||
"priority": "medium"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Example 5: Handling blockers (keeping task in_progress)**
|
||||
```json
|
||||
{
|
||||
"action": "write",
|
||||
"todos": [
|
||||
{
|
||||
"id": "1",
|
||||
"content": "Deploy to production",
|
||||
"status": "in_progress",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "2",
|
||||
"content": "BLOCKER: Fix failing deployment pipeline",
|
||||
"status": "pending",
|
||||
"priority": "high"
|
||||
},
|
||||
{
|
||||
"id": "3",
|
||||
"content": "Update documentation",
|
||||
"status": "pending",
|
||||
"priority": "low"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Common Scenarios
|
||||
|
||||
**Multi-file refactoring:** Create todos for each file that needs updating
|
||||
**Performance optimization:** List specific bottlenecks as individual tasks
|
||||
**Bug fixing:** Track reproduction, diagnosis, fix, and verification as separate tasks
|
||||
**Feature implementation:** Break down into UI, logic, tests, and documentation tasks
|
||||
|
||||
Remember: When writing, you must include ALL todos you want to keep. Any todo not in the list will be removed. Be proactive with task management to demonstrate thoroughness and ensure all requirements are completed successfully.
|
||||
42
vibe/core/tools/builtins/prompts/write_file.md
Normal file
42
vibe/core/tools/builtins/prompts/write_file.md
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
Use `write_file` to write content to a file.
|
||||
|
||||
**Arguments:**
|
||||
- `path`: The file path (relative or absolute)
|
||||
- `content`: The content to write to the file
|
||||
- `overwrite`: Must be set to `true` to overwrite an existing file (default: `false`)
|
||||
|
||||
**IMPORTANT SAFETY RULES:**
|
||||
|
||||
- By default, the tool will **fail if the file already exists** to prevent accidental data loss
|
||||
- To **overwrite** an existing file, you **MUST** set `overwrite: true`
|
||||
- To **create a new file**, just provide the `path` and `content` (overwrite defaults to false)
|
||||
- If parent directories don't exist, they will be created automatically
|
||||
|
||||
**BEST PRACTICES:**
|
||||
|
||||
- **ALWAYS** use the `read_file` tool first before overwriting an existing file to understand its current contents
|
||||
- **ALWAYS** prefer using `search_replace` to edit existing files rather than overwriting them completely
|
||||
- **NEVER** write new files unless explicitly required - prefer modifying existing files
|
||||
- **NEVER** proactively create documentation files (*.md) or README files unless explicitly requested
|
||||
- **AVOID** using emojis in file content unless the user explicitly requests them
|
||||
|
||||
**Usage Examples:**
|
||||
|
||||
```python
|
||||
# Create a new file (will error if file exists)
|
||||
write_file(
|
||||
path="src/new_module.py",
|
||||
content="def hello():\n return 'Hello World'"
|
||||
)
|
||||
|
||||
# Overwrite an existing file (must read it first!)
|
||||
# First: read_file(path="src/existing.py")
|
||||
# Then:
|
||||
write_file(
|
||||
path="src/existing.py",
|
||||
content="# Updated content\ndef new_function():\n pass",
|
||||
overwrite=True
|
||||
)
|
||||
```
|
||||
|
||||
**Remember:** For editing existing files, prefer `search_replace` over `write_file` to preserve unchanged portions and avoid accidental data loss.
|
||||
231
vibe/core/tools/builtins/read_file.py
Normal file
231
vibe/core/tools/builtins/read_file.py
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, ClassVar, NamedTuple, final
|
||||
|
||||
import aiofiles
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
BaseToolConfig,
|
||||
BaseToolState,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
)
|
||||
from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class _ReadResult(NamedTuple):
|
||||
lines: list[str]
|
||||
bytes_read: int
|
||||
was_truncated: bool
|
||||
|
||||
|
||||
class ReadFileArgs(BaseModel):
|
||||
path: str
|
||||
offset: int = Field(
|
||||
default=0,
|
||||
description="Line number to start reading from (0-indexed, inclusive).",
|
||||
)
|
||||
limit: int | None = Field(
|
||||
default=None, description="Maximum number of lines to read."
|
||||
)
|
||||
|
||||
|
||||
class ReadFileResult(BaseModel):
|
||||
path: str
|
||||
content: str
|
||||
lines_read: int
|
||||
was_truncated: bool = Field(
|
||||
description="True if the reading was stopped due to the max_read_bytes limit."
|
||||
)
|
||||
|
||||
|
||||
class ReadFileToolConfig(BaseToolConfig):
|
||||
permission: ToolPermission = ToolPermission.ALWAYS
|
||||
|
||||
max_read_bytes: int = Field(
|
||||
default=64_000, description="Maximum total bytes to read from a file in one go."
|
||||
)
|
||||
max_state_history: int = Field(
|
||||
default=10, description="Number of recently read files to remember in state."
|
||||
)
|
||||
|
||||
|
||||
class ReadFileState(BaseToolState):
|
||||
recently_read_files: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ReadFile(
|
||||
BaseTool[ReadFileArgs, ReadFileResult, ReadFileToolConfig, ReadFileState],
|
||||
ToolUIData[ReadFileArgs, ReadFileResult],
|
||||
):
|
||||
description: ClassVar[str] = (
|
||||
"Read a UTF-8 file, returning content from a specific line range. "
|
||||
"Reading is capped by a byte limit for safety."
|
||||
)
|
||||
|
||||
@final
|
||||
async def run(self, args: ReadFileArgs) -> ReadFileResult:
|
||||
file_path = self._prepare_and_validate_path(args)
|
||||
|
||||
read_result = await self._read_file(args, file_path)
|
||||
|
||||
self._update_state_history(file_path)
|
||||
|
||||
return ReadFileResult(
|
||||
path=str(file_path),
|
||||
content="".join(read_result.lines),
|
||||
lines_read=len(read_result.lines),
|
||||
was_truncated=read_result.was_truncated,
|
||||
)
|
||||
|
||||
def check_allowlist_denylist(self, args: ReadFileArgs) -> ToolPermission | None:
|
||||
import fnmatch
|
||||
|
||||
file_path = Path(args.path).expanduser()
|
||||
if not file_path.is_absolute():
|
||||
file_path = self.config.effective_workdir / file_path
|
||||
file_str = str(file_path)
|
||||
|
||||
for pattern in self.config.denylist:
|
||||
if fnmatch.fnmatch(file_str, pattern):
|
||||
return ToolPermission.NEVER
|
||||
|
||||
for pattern in self.config.allowlist:
|
||||
if fnmatch.fnmatch(file_str, pattern):
|
||||
return ToolPermission.ALWAYS
|
||||
|
||||
return None
|
||||
|
||||
def _prepare_and_validate_path(self, args: ReadFileArgs) -> Path:
|
||||
self._validate_inputs(args)
|
||||
|
||||
file_path = Path(args.path).expanduser()
|
||||
if not file_path.is_absolute():
|
||||
file_path = self.config.effective_workdir / file_path
|
||||
|
||||
self._validate_path(file_path)
|
||||
return file_path
|
||||
|
||||
async def _read_file(self, args: ReadFileArgs, file_path: Path) -> _ReadResult:
|
||||
try:
|
||||
lines_to_return: list[str] = []
|
||||
bytes_read = 0
|
||||
was_truncated = False
|
||||
|
||||
async with aiofiles.open(file_path, encoding="utf-8", errors="ignore") as f:
|
||||
line_index = 0
|
||||
async for line in f:
|
||||
if line_index < args.offset:
|
||||
line_index += 1
|
||||
continue
|
||||
|
||||
if args.limit is not None and len(lines_to_return) >= args.limit:
|
||||
break
|
||||
|
||||
line_bytes = len(line.encode("utf-8"))
|
||||
if bytes_read + line_bytes > self.config.max_read_bytes:
|
||||
was_truncated = True
|
||||
break
|
||||
|
||||
lines_to_return.append(line)
|
||||
bytes_read += line_bytes
|
||||
line_index += 1
|
||||
|
||||
return _ReadResult(
|
||||
lines=lines_to_return,
|
||||
bytes_read=bytes_read,
|
||||
was_truncated=was_truncated,
|
||||
)
|
||||
|
||||
except OSError as exc:
|
||||
raise ToolError(f"Error reading {file_path}: {exc}") from exc
|
||||
|
||||
def _validate_inputs(self, args: ReadFileArgs) -> None:
|
||||
if not args.path.strip():
|
||||
raise ToolError("Path cannot be empty")
|
||||
if args.offset < 0:
|
||||
raise ToolError("Offset cannot be negative")
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
raise ToolError("Limit, if provided, must be a positive number")
|
||||
|
||||
def _validate_path(self, file_path: Path) -> None:
|
||||
try:
|
||||
resolved_path = file_path.resolve()
|
||||
except ValueError:
|
||||
raise ToolError(
|
||||
f"Security error: Cannot read path '{file_path}' outside of the project directory '{self.config.effective_workdir}'."
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise ToolError(f"File not found at: {file_path}")
|
||||
|
||||
if not resolved_path.exists():
|
||||
raise ToolError(f"File not found at: {file_path}")
|
||||
if resolved_path.is_dir():
|
||||
raise ToolError(f"Path is a directory, not a file: {file_path}")
|
||||
|
||||
def _update_state_history(self, file_path: Path) -> None:
|
||||
self.state.recently_read_files.append(str(file_path.resolve()))
|
||||
if len(self.state.recently_read_files) > self.config.max_state_history:
|
||||
self.state.recently_read_files.pop(0)
|
||||
|
||||
@classmethod
|
||||
def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
|
||||
if not isinstance(event.args, ReadFileArgs):
|
||||
return ToolCallDisplay(summary="read_file")
|
||||
|
||||
summary = f"read_file: {event.args.path}"
|
||||
if event.args.offset > 0 or event.args.limit is not None:
|
||||
parts = []
|
||||
if event.args.offset > 0:
|
||||
parts.append(f"from line {event.args.offset}")
|
||||
if event.args.limit is not None:
|
||||
parts.append(f"limit {event.args.limit} lines")
|
||||
summary += f" ({', '.join(parts)})"
|
||||
|
||||
return ToolCallDisplay(
|
||||
summary=summary,
|
||||
details={
|
||||
"path": event.args.path,
|
||||
"offset": event.args.offset,
|
||||
"limit": event.args.limit,
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
|
||||
if not isinstance(event.result, ReadFileResult):
|
||||
return ToolResultDisplay(
|
||||
success=False, message=event.error or event.skip_reason or "No result"
|
||||
)
|
||||
|
||||
path_obj = Path(event.result.path)
|
||||
message = f"Read {event.result.lines_read} line{'' if event.result.lines_read <= 1 else 's'} from {path_obj.name}"
|
||||
if event.result.was_truncated:
|
||||
message += " (truncated)"
|
||||
|
||||
return ToolResultDisplay(
|
||||
success=True,
|
||||
message=message,
|
||||
warnings=["File was truncated due to size limit"]
|
||||
if event.result.was_truncated
|
||||
else [],
|
||||
details={
|
||||
"path": str(event.result.path),
|
||||
"lines_read": event.result.lines_read,
|
||||
"was_truncated": event.result.was_truncated,
|
||||
"content": event.result.content,
|
||||
"file_extension": path_obj.suffix.lstrip(".")
|
||||
if path_obj.suffix
|
||||
else "text",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_status_text(cls) -> str:
|
||||
return "Reading file"
|
||||
454
vibe/core/tools/builtins/search_replace.py
Normal file
454
vibe/core/tools/builtins/search_replace.py
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import difflib
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
from typing import ClassVar, NamedTuple, final
|
||||
|
||||
import aiofiles
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import BaseTool, BaseToolConfig, BaseToolState, ToolError
|
||||
from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
_BLOCK_RE = re.compile(
|
||||
r"<{5,} SEARCH\r?\n(.*?)\r?\n?={5,}\r?\n(.*?)\r?\n?>{5,} REPLACE", flags=re.DOTALL
|
||||
)
|
||||
|
||||
_BLOCK_WITH_FENCE_RE = re.compile(
|
||||
r"```[\s\S]*?\n<{5,} SEARCH\r?\n(.*?)\r?\n?={5,}\r?\n(.*?)\r?\n?>{5,} REPLACE\s*\n```",
|
||||
flags=re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
class SearchReplaceBlock(NamedTuple):
|
||||
search: str
|
||||
replace: str
|
||||
|
||||
|
||||
class FuzzyMatch(NamedTuple):
|
||||
similarity: float
|
||||
start_line: int
|
||||
end_line: int
|
||||
text: str
|
||||
|
||||
|
||||
class BlockApplyResult(NamedTuple):
|
||||
content: str
|
||||
applied: int
|
||||
errors: list[str]
|
||||
warnings: list[str]
|
||||
|
||||
|
||||
class SearchReplaceArgs(BaseModel):
|
||||
file_path: str
|
||||
content: str
|
||||
|
||||
|
||||
class SearchReplaceResult(BaseModel):
|
||||
file: str
|
||||
blocks_applied: int
|
||||
lines_changed: int
|
||||
content: str
|
||||
warnings: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class SearchReplaceConfig(BaseToolConfig):
|
||||
max_content_size: int = 100_000
|
||||
create_backup: bool = False
|
||||
fuzzy_threshold: float = 0.9
|
||||
|
||||
|
||||
class SearchReplaceState(BaseToolState):
|
||||
pass
|
||||
|
||||
|
||||
class SearchReplace(
|
||||
BaseTool[
|
||||
SearchReplaceArgs, SearchReplaceResult, SearchReplaceConfig, SearchReplaceState
|
||||
],
|
||||
ToolUIData[SearchReplaceArgs, SearchReplaceResult],
|
||||
):
|
||||
description: ClassVar[str] = (
|
||||
"Replace sections of files using SEARCH/REPLACE blocks. "
|
||||
"Supports fuzzy matching and detailed error reporting. "
|
||||
"Format: <<<<<<< SEARCH\\n[text]\\n=======\\n[replacement]\\n>>>>>>> REPLACE"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
|
||||
if not isinstance(event.args, SearchReplaceArgs):
|
||||
return ToolCallDisplay(summary="Invalid arguments")
|
||||
|
||||
args = event.args
|
||||
blocks = cls._parse_search_replace_blocks(args.content)
|
||||
|
||||
return ToolCallDisplay(
|
||||
summary=f"Patching {args.file_path} ({len(blocks)} blocks)",
|
||||
content=args.content,
|
||||
details={
|
||||
"path": args.file_path,
|
||||
"blocks_count": len(blocks),
|
||||
"original_path": args.file_path,
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
|
||||
if isinstance(event.result, SearchReplaceResult):
|
||||
return ToolResultDisplay(
|
||||
success=True,
|
||||
message=f"Applied {event.result.blocks_applied} blocks",
|
||||
warnings=event.result.warnings,
|
||||
details={
|
||||
"lines_changed": event.result.lines_changed,
|
||||
"content": event.result.content,
|
||||
},
|
||||
)
|
||||
|
||||
return ToolResultDisplay(success=True, message="Patch applied")
|
||||
|
||||
@classmethod
|
||||
def get_status_text(cls) -> str:
|
||||
return "Editing files"
|
||||
|
||||
@final
|
||||
async def run(self, args: SearchReplaceArgs) -> SearchReplaceResult:
|
||||
file_path, search_replace_blocks = self._prepare_and_validate_args(args)
|
||||
|
||||
original_content = await self._read_file(file_path)
|
||||
|
||||
block_result = self._apply_blocks(
|
||||
original_content,
|
||||
search_replace_blocks,
|
||||
file_path,
|
||||
self.config.fuzzy_threshold,
|
||||
)
|
||||
|
||||
if block_result.errors:
|
||||
error_message = "SEARCH/REPLACE blocks failed:\n" + "\n\n".join(
|
||||
block_result.errors
|
||||
)
|
||||
if block_result.warnings:
|
||||
error_message += "\n\nWarnings encountered:\n" + "\n".join(
|
||||
block_result.warnings
|
||||
)
|
||||
raise ToolError(error_message)
|
||||
|
||||
modified_content = block_result.content
|
||||
|
||||
# Calculate line changes
|
||||
if modified_content == original_content:
|
||||
lines_changed = 0
|
||||
else:
|
||||
original_lines = len(original_content.splitlines())
|
||||
new_lines = len(modified_content.splitlines())
|
||||
lines_changed = new_lines - original_lines
|
||||
|
||||
try:
|
||||
if self.config.create_backup:
|
||||
await self._backup_file(file_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
await self._write_file(file_path, modified_content)
|
||||
|
||||
return SearchReplaceResult(
|
||||
file=str(file_path),
|
||||
blocks_applied=block_result.applied,
|
||||
lines_changed=lines_changed,
|
||||
warnings=block_result.warnings,
|
||||
content=args.content,
|
||||
)
|
||||
|
||||
@final
|
||||
def _prepare_and_validate_args(
|
||||
self, args: SearchReplaceArgs
|
||||
) -> tuple[Path, list[SearchReplaceBlock]]:
|
||||
file_path_str = args.file_path.strip()
|
||||
content = args.content.strip()
|
||||
|
||||
if not file_path_str:
|
||||
raise ToolError("File path cannot be empty")
|
||||
|
||||
if len(content) > self.config.max_content_size:
|
||||
raise ToolError(
|
||||
f"Content size ({len(content)} bytes) exceeds max_content_size "
|
||||
f"({self.config.max_content_size} bytes)"
|
||||
)
|
||||
|
||||
if not content:
|
||||
raise ToolError("Empty content provided")
|
||||
|
||||
project_root = self.config.effective_workdir
|
||||
file_path = Path(file_path_str).expanduser()
|
||||
if not file_path.is_absolute():
|
||||
file_path = project_root / file_path
|
||||
file_path = file_path.resolve()
|
||||
|
||||
if not file_path.exists():
|
||||
raise ToolError(f"File does not exist: {file_path}")
|
||||
|
||||
if not file_path.is_file():
|
||||
raise ToolError(f"Path is not a file: {file_path}")
|
||||
|
||||
search_replace_blocks = self._parse_search_replace_blocks(content)
|
||||
if not search_replace_blocks:
|
||||
raise ToolError(
|
||||
"No valid SEARCH/REPLACE blocks found in content.\n"
|
||||
"Expected format:\n"
|
||||
"<<<<<<< SEARCH\n"
|
||||
"[exact content to find]\n"
|
||||
"=======\n"
|
||||
"[new content to replace with]\n"
|
||||
">>>>>>> REPLACE"
|
||||
)
|
||||
|
||||
return file_path, search_replace_blocks
|
||||
|
||||
async def _read_file(self, file_path: Path) -> str:
|
||||
try:
|
||||
async with aiofiles.open(file_path, encoding="utf-8") as f:
|
||||
return await f.read()
|
||||
except UnicodeDecodeError as e:
|
||||
raise ToolError(f"Unicode decode error reading {file_path}: {e}") from e
|
||||
except PermissionError:
|
||||
raise ToolError(f"Permission denied reading file: {file_path}")
|
||||
except Exception as e:
|
||||
raise ToolError(f"Unexpected error reading {file_path}: {e}") from e
|
||||
|
||||
async def _backup_file(self, file_path: Path) -> None:
|
||||
shutil.copy2(file_path, file_path.with_suffix(file_path.suffix + ".bak"))
|
||||
|
||||
async def _write_file(self, file_path: Path, content: str) -> None:
|
||||
try:
|
||||
async with aiofiles.open(file_path, mode="w", encoding="utf-8") as f:
|
||||
await f.write(content)
|
||||
except PermissionError:
|
||||
raise ToolError(f"Permission denied writing to file: {file_path}")
|
||||
except OSError as e:
|
||||
raise ToolError(f"OS error writing to {file_path}: {e}") from e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Unexpected error writing to {file_path}: {e}") from e
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _apply_blocks(
|
||||
content: str,
|
||||
blocks: list[SearchReplaceBlock],
|
||||
filepath: Path,
|
||||
fuzzy_threshold: float = 0.9,
|
||||
) -> BlockApplyResult:
|
||||
applied = 0
|
||||
errors: list[str] = []
|
||||
warnings: list[str] = []
|
||||
current_content = content
|
||||
|
||||
for i, (search, replace) in enumerate(blocks, 1):
|
||||
if search not in current_content:
|
||||
context = SearchReplace._find_search_context(current_content, search)
|
||||
fuzzy_context = SearchReplace._find_fuzzy_match_context(
|
||||
current_content, search, fuzzy_threshold
|
||||
)
|
||||
|
||||
error_msg = (
|
||||
f"SEARCH/REPLACE block {i} failed: Search text not found in {filepath}\n"
|
||||
f"Search text was:\n{search!r}\n"
|
||||
f"Context analysis:\n{context}"
|
||||
)
|
||||
|
||||
if fuzzy_context:
|
||||
error_msg += f"\n{fuzzy_context}"
|
||||
|
||||
error_msg += (
|
||||
"\nDebugging tips:\n"
|
||||
"1. Check for exact whitespace/indentation match\n"
|
||||
"2. Verify line endings match the file exactly (\\r\\n vs \\n)\n"
|
||||
"3. Ensure the search text hasn't been modified by previous blocks or user edits\n"
|
||||
"4. Check for typos or case sensitivity issues"
|
||||
)
|
||||
|
||||
errors.append(error_msg)
|
||||
continue
|
||||
|
||||
occurrences = current_content.count(search)
|
||||
if occurrences > 1:
|
||||
warning_msg = (
|
||||
f"Search text in block {i} appears {occurrences} times in the file. "
|
||||
f"Only the first occurrence will be replaced. Consider making your "
|
||||
f"search pattern more specific to avoid unintended changes."
|
||||
)
|
||||
warnings.append(warning_msg)
|
||||
|
||||
current_content = current_content.replace(search, replace, 1)
|
||||
applied += 1
|
||||
|
||||
return BlockApplyResult(
|
||||
content=current_content, applied=applied, errors=errors, warnings=warnings
|
||||
)
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _find_fuzzy_match_context(
|
||||
content: str, search_text: str, threshold: float = 0.9
|
||||
) -> str | None:
|
||||
best_match = SearchReplace._find_best_fuzzy_match(
|
||||
content, search_text, threshold
|
||||
)
|
||||
|
||||
if not best_match:
|
||||
return None
|
||||
|
||||
diff = SearchReplace._create_unified_diff(
|
||||
search_text, best_match.text, "SEARCH", "CLOSEST MATCH"
|
||||
)
|
||||
|
||||
similarity_pct = best_match.similarity * 100
|
||||
|
||||
return (
|
||||
f"Closest fuzzy match (similarity {similarity_pct:.1f}%) "
|
||||
f"at lines {best_match.start_line}–{best_match.end_line}:\n"
|
||||
f"```diff\n{diff}\n```"
|
||||
)
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _find_best_fuzzy_match( # noqa: PLR0914
|
||||
content: str, search_text: str, threshold: float = 0.9
|
||||
) -> FuzzyMatch | None:
|
||||
content_lines = content.split("\n")
|
||||
search_lines = search_text.split("\n")
|
||||
window_size = len(search_lines)
|
||||
|
||||
if window_size == 0:
|
||||
return None
|
||||
|
||||
non_empty_search = [line for line in search_lines if line.strip()]
|
||||
if not non_empty_search:
|
||||
return None
|
||||
|
||||
first_anchor = non_empty_search[0]
|
||||
last_anchor = (
|
||||
non_empty_search[-1] if len(non_empty_search) > 1 else first_anchor
|
||||
)
|
||||
|
||||
candidate_starts = set()
|
||||
spread = 5
|
||||
|
||||
for i, line in enumerate(content_lines):
|
||||
if first_anchor in line or last_anchor in line:
|
||||
start_min = max(0, i - spread)
|
||||
start_max = min(len(content_lines) - window_size + 1, i + spread + 1)
|
||||
for s in range(start_min, start_max):
|
||||
candidate_starts.add(s)
|
||||
|
||||
if not candidate_starts:
|
||||
max_positions = min(len(content_lines) - window_size + 1, 100)
|
||||
candidate_starts = set(range(0, max_positions))
|
||||
|
||||
best_match = None
|
||||
best_similarity = 0.0
|
||||
|
||||
for start in candidate_starts:
|
||||
end = start + window_size
|
||||
window_text = "\n".join(content_lines[start:end])
|
||||
|
||||
matcher = difflib.SequenceMatcher(None, search_text, window_text)
|
||||
similarity = matcher.ratio()
|
||||
|
||||
if similarity >= threshold and similarity > best_similarity:
|
||||
best_similarity = similarity
|
||||
best_match = FuzzyMatch(
|
||||
similarity=similarity,
|
||||
start_line=start + 1, # 1-based line numbers
|
||||
end_line=end,
|
||||
text=window_text,
|
||||
)
|
||||
|
||||
return best_match
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _create_unified_diff(
|
||||
text1: str, text2: str, label1: str = "SEARCH", label2: str = "CLOSEST MATCH"
|
||||
) -> str:
|
||||
lines1 = text1.splitlines(keepends=True)
|
||||
lines2 = text2.splitlines(keepends=True)
|
||||
|
||||
lines1 = [line if line.endswith("\n") else line + "\n" for line in lines1]
|
||||
lines2 = [line if line.endswith("\n") else line + "\n" for line in lines2]
|
||||
|
||||
diff = difflib.unified_diff(
|
||||
lines1, lines2, fromfile=label1, tofile=label2, lineterm="", n=3
|
||||
)
|
||||
|
||||
diff_lines = list(diff)
|
||||
|
||||
if diff_lines and not diff_lines[0].startswith("==="):
|
||||
diff_lines.insert(2, "=" * 67 + "\n")
|
||||
|
||||
result = "".join(diff_lines)
|
||||
|
||||
max_chars = 2000
|
||||
if len(result) > max_chars:
|
||||
result = result[:max_chars] + "\n...(diff truncated)"
|
||||
|
||||
return result.rstrip()
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _parse_search_replace_blocks(content: str) -> list[SearchReplaceBlock]:
|
||||
"""Parse SEARCH/REPLACE blocks from content.
|
||||
|
||||
Supports two formats:
|
||||
1. With code block fences (```...```)
|
||||
2. Without code block fences
|
||||
"""
|
||||
matches = _BLOCK_WITH_FENCE_RE.findall(content)
|
||||
|
||||
if not matches:
|
||||
matches = _BLOCK_RE.findall(content)
|
||||
|
||||
return [
|
||||
SearchReplaceBlock(
|
||||
search=search.rstrip("\r\n"), replace=replace.rstrip("\r\n")
|
||||
)
|
||||
for search, replace in matches
|
||||
]
|
||||
|
||||
@final
|
||||
@staticmethod
|
||||
def _find_search_context(
|
||||
content: str, search_text: str, max_context: int = 5
|
||||
) -> str:
|
||||
lines = content.split("\n")
|
||||
search_lines = search_text.split("\n")
|
||||
|
||||
if not search_lines:
|
||||
return "Search text is empty"
|
||||
|
||||
first_search_line = search_lines[0].strip()
|
||||
if not first_search_line:
|
||||
return "First line of search text is empty or whitespace only"
|
||||
|
||||
matches = []
|
||||
for i, line in enumerate(lines):
|
||||
if first_search_line in line:
|
||||
matches.append(i)
|
||||
|
||||
if not matches:
|
||||
return f"First search line '{first_search_line}' not found anywhere in file"
|
||||
|
||||
context_lines = []
|
||||
for match_idx in matches[:3]:
|
||||
start = max(0, match_idx - max_context)
|
||||
end = min(len(lines), match_idx + max_context + 1)
|
||||
|
||||
context_lines.append(f"\nPotential match area around line {match_idx + 1}:")
|
||||
for i in range(start, end):
|
||||
marker = ">>>" if i == match_idx else " "
|
||||
context_lines.append(f"{marker} {i + 1:3d}: {lines[i]}")
|
||||
|
||||
return "\n".join(context_lines)
|
||||
147
vibe/core/tools/builtins/todo.py
Normal file
147
vibe/core/tools/builtins/todo.py
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum, auto
|
||||
from typing import ClassVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
BaseToolConfig,
|
||||
BaseToolState,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
)
|
||||
from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class TodoStatus(StrEnum):
|
||||
PENDING = auto()
|
||||
IN_PROGRESS = auto()
|
||||
COMPLETED = auto()
|
||||
CANCELLED = auto()
|
||||
|
||||
|
||||
class TodoPriority(StrEnum):
|
||||
LOW = auto()
|
||||
MEDIUM = auto()
|
||||
HIGH = auto()
|
||||
|
||||
|
||||
class TodoItem(BaseModel):
|
||||
id: str
|
||||
content: str
|
||||
status: TodoStatus = TodoStatus.PENDING
|
||||
priority: TodoPriority = TodoPriority.MEDIUM
|
||||
|
||||
|
||||
class TodoArgs(BaseModel):
|
||||
action: str = Field(description="Either 'read' or 'write'")
|
||||
todos: list[TodoItem] | None = Field(
|
||||
default=None, description="Complete list of todos when writing."
|
||||
)
|
||||
|
||||
|
||||
class TodoResult(BaseModel):
|
||||
message: str
|
||||
todos: list[TodoItem]
|
||||
total_count: int
|
||||
|
||||
|
||||
class TodoConfig(BaseToolConfig):
|
||||
permission: ToolPermission = ToolPermission.ALWAYS
|
||||
max_todos: int = 100
|
||||
|
||||
|
||||
class TodoState(BaseToolState):
|
||||
todos: list[TodoItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Todo(
|
||||
BaseTool[TodoArgs, TodoResult, TodoConfig, TodoState],
|
||||
ToolUIData[TodoArgs, TodoResult],
|
||||
):
|
||||
description: ClassVar[str] = (
|
||||
"Manage todos. Use action='read' to view, action='write' with complete list to update."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
|
||||
if not isinstance(event.args, TodoArgs):
|
||||
return ToolCallDisplay(summary="Invalid arguments")
|
||||
|
||||
args = event.args
|
||||
|
||||
match args.action:
|
||||
case "read":
|
||||
return ToolCallDisplay(
|
||||
summary="Reading todos", details={"action": "read"}
|
||||
)
|
||||
case "write":
|
||||
count = len(args.todos) if args.todos else 0
|
||||
return ToolCallDisplay(
|
||||
summary=f"Writing {count} todos",
|
||||
details={"action": "write", "count": count},
|
||||
)
|
||||
case _:
|
||||
return ToolCallDisplay(
|
||||
summary=f"Unknown action: {args.action}",
|
||||
details={"action": args.action},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
|
||||
if not isinstance(event.result, TodoResult):
|
||||
return ToolResultDisplay(success=True, message="Success")
|
||||
|
||||
result = event.result
|
||||
|
||||
by_status = {"in_progress": [], "pending": [], "completed": [], "cancelled": []}
|
||||
|
||||
for todo in result.todos:
|
||||
by_status[todo.status].append({"content": todo.content, "id": todo.id})
|
||||
|
||||
return ToolResultDisplay(
|
||||
success=True,
|
||||
message=result.message,
|
||||
details={"todos_by_status": by_status, "total_count": result.total_count},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_status_text(cls) -> str:
|
||||
return "Managing todos"
|
||||
|
||||
async def run(self, args: TodoArgs) -> TodoResult:
|
||||
match args.action:
|
||||
case "read":
|
||||
return self._read_todos()
|
||||
case "write":
|
||||
return self._write_todos(args.todos or [])
|
||||
case _:
|
||||
raise ToolError(
|
||||
f"Invalid action '{args.action}'. Use 'read' or 'write'."
|
||||
)
|
||||
|
||||
def _read_todos(self) -> TodoResult:
|
||||
return TodoResult(
|
||||
message=f"Retrieved {len(self.state.todos)} todos",
|
||||
todos=self.state.todos,
|
||||
total_count=len(self.state.todos),
|
||||
)
|
||||
|
||||
def _write_todos(self, todos: list[TodoItem]) -> TodoResult:
|
||||
if len(todos) > self.config.max_todos:
|
||||
raise ToolError(f"Cannot store more than {self.config.max_todos} todos")
|
||||
|
||||
ids = [todo.id for todo in todos]
|
||||
if len(ids) != len(set(ids)):
|
||||
raise ToolError("Todo IDs must be unique")
|
||||
|
||||
self.state.todos = todos
|
||||
|
||||
return TodoResult(
|
||||
message=f"Updated {len(todos)} todos",
|
||||
todos=self.state.todos,
|
||||
total_count=len(self.state.todos),
|
||||
)
|
||||
167
vibe/core/tools/builtins/write_file.py
Normal file
167
vibe/core/tools/builtins/write_file.py
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import ClassVar, final
|
||||
|
||||
import aiofiles
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from vibe.core.tools.base import (
|
||||
BaseTool,
|
||||
BaseToolConfig,
|
||||
BaseToolState,
|
||||
ToolError,
|
||||
ToolPermission,
|
||||
)
|
||||
from vibe.core.tools.ui import ToolCallDisplay, ToolResultDisplay, ToolUIData
|
||||
from vibe.core.types import ToolCallEvent, ToolResultEvent
|
||||
|
||||
|
||||
class WriteFileArgs(BaseModel):
|
||||
path: str
|
||||
content: str
|
||||
overwrite: bool = Field(
|
||||
default=False, description="Must be set to true to overwrite an existing file."
|
||||
)
|
||||
|
||||
|
||||
class WriteFileResult(BaseModel):
|
||||
path: str
|
||||
bytes_written: int
|
||||
file_existed: bool
|
||||
content: str
|
||||
|
||||
|
||||
class WriteFileConfig(BaseToolConfig):
|
||||
permission: ToolPermission = ToolPermission.ASK
|
||||
max_write_bytes: int = 64_000
|
||||
create_parent_dirs: bool = True
|
||||
|
||||
|
||||
class WriteFileState(BaseToolState):
|
||||
recently_written_files: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class WriteFile(
|
||||
BaseTool[WriteFileArgs, WriteFileResult, WriteFileConfig, WriteFileState],
|
||||
ToolUIData[WriteFileArgs, WriteFileResult],
|
||||
):
|
||||
description: ClassVar[str] = (
|
||||
"Create or overwrite a UTF-8 file. Fails if file exists unless 'overwrite=True'."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_call_display(cls, event: ToolCallEvent) -> ToolCallDisplay:
|
||||
if not isinstance(event.args, WriteFileArgs):
|
||||
return ToolCallDisplay(summary="Invalid arguments")
|
||||
|
||||
args = event.args
|
||||
file_ext = Path(args.path).suffix.lstrip(".")
|
||||
|
||||
return ToolCallDisplay(
|
||||
summary=f"Writing {args.path}{' (overwrite)' if args.overwrite else ''}",
|
||||
content=args.content,
|
||||
details={
|
||||
"path": args.path,
|
||||
"overwrite": args.overwrite,
|
||||
"file_extension": file_ext,
|
||||
"content": args.content,
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_result_display(cls, event: ToolResultEvent) -> ToolResultDisplay:
|
||||
if isinstance(event.result, WriteFileResult):
|
||||
action = "Overwritten" if event.result.file_existed else "Created"
|
||||
return ToolResultDisplay(
|
||||
success=True,
|
||||
message=f"{action} {Path(event.result.path).name}",
|
||||
details={
|
||||
"bytes_written": event.result.bytes_written,
|
||||
"path": event.result.path,
|
||||
"content": event.result.content,
|
||||
},
|
||||
)
|
||||
|
||||
return ToolResultDisplay(success=True, message="File written")
|
||||
|
||||
@classmethod
|
||||
def get_status_text(cls) -> str:
|
||||
return "Writing file"
|
||||
|
||||
def check_allowlist_denylist(self, args: WriteFileArgs) -> ToolPermission | None:
|
||||
import fnmatch
|
||||
|
||||
file_path = Path(args.path).expanduser()
|
||||
if not file_path.is_absolute():
|
||||
file_path = self.config.effective_workdir / file_path
|
||||
file_str = str(file_path)
|
||||
|
||||
for pattern in self.config.denylist:
|
||||
if fnmatch.fnmatch(file_str, pattern):
|
||||
return ToolPermission.NEVER
|
||||
|
||||
for pattern in self.config.allowlist:
|
||||
if fnmatch.fnmatch(file_str, pattern):
|
||||
return ToolPermission.ALWAYS
|
||||
|
||||
return None
|
||||
|
||||
@final
|
||||
async def run(self, args: WriteFileArgs) -> WriteFileResult:
|
||||
file_path, file_existed, content_bytes = self._prepare_and_validate_path(args)
|
||||
|
||||
await self._write_file(args, file_path)
|
||||
|
||||
BUFFER_SIZE = 10
|
||||
self.state.recently_written_files.append(str(file_path))
|
||||
if len(self.state.recently_written_files) > BUFFER_SIZE:
|
||||
self.state.recently_written_files.pop(0)
|
||||
|
||||
return WriteFileResult(
|
||||
path=str(file_path),
|
||||
bytes_written=content_bytes,
|
||||
file_existed=file_existed,
|
||||
content=args.content,
|
||||
)
|
||||
|
||||
def _prepare_and_validate_path(self, args: WriteFileArgs) -> tuple[Path, bool, int]:
|
||||
if not args.path.strip():
|
||||
raise ToolError("Path cannot be empty")
|
||||
|
||||
content_bytes = len(args.content.encode("utf-8"))
|
||||
if content_bytes > self.config.max_write_bytes:
|
||||
raise ToolError(
|
||||
f"Content exceeds {self.config.max_write_bytes} bytes limit"
|
||||
)
|
||||
|
||||
file_path = Path(args.path).expanduser()
|
||||
if not file_path.is_absolute():
|
||||
file_path = self.config.effective_workdir / file_path
|
||||
file_path = file_path.resolve()
|
||||
|
||||
try:
|
||||
file_path.relative_to(self.config.effective_workdir.resolve())
|
||||
except ValueError:
|
||||
raise ToolError(f"Cannot write outside project directory: {file_path}")
|
||||
|
||||
file_existed = file_path.exists()
|
||||
|
||||
if file_existed and not args.overwrite:
|
||||
raise ToolError(
|
||||
f"File '{file_path}' exists. Set overwrite=True to replace."
|
||||
)
|
||||
|
||||
if self.config.create_parent_dirs:
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
elif not file_path.parent.exists():
|
||||
raise ToolError(f"Parent directory does not exist: {file_path.parent}")
|
||||
|
||||
return file_path, file_existed, content_bytes
|
||||
|
||||
async def _write_file(self, args: WriteFileArgs, file_path: Path) -> None:
|
||||
try:
|
||||
async with aiofiles.open(file_path, mode="w", encoding="utf-8") as f:
|
||||
await f.write(args.content)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error writing {file_path}: {e}") from e
|
||||
273
vibe/core/tools/manager.py
Normal file
273
vibe/core/tools/manager.py
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
import importlib.util
|
||||
import inspect
|
||||
from logging import getLogger
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from vibe import VIBE_ROOT
|
||||
from vibe.core.config import get_vibe_home
|
||||
from vibe.core.tools.base import BaseTool, BaseToolConfig
|
||||
from vibe.core.tools.mcp import (
|
||||
RemoteTool,
|
||||
create_mcp_http_proxy_tool_class,
|
||||
create_mcp_stdio_proxy_tool_class,
|
||||
list_tools_http,
|
||||
list_tools_stdio,
|
||||
)
|
||||
from vibe.core.utils import run_sync
|
||||
|
||||
logger = getLogger("vibe")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import MCPHttp, MCPStdio, MCPStreamableHttp, VibeConfig
|
||||
|
||||
|
||||
class NoSuchToolError(Exception):
|
||||
"""Exception raised when a tool is not found."""
|
||||
|
||||
|
||||
DEFAULT_TOOL_DIR = VIBE_ROOT / "core" / "tools" / "builtins"
|
||||
|
||||
|
||||
class ToolManager:
|
||||
"""Manages tool discovery and instantiation for an Agent.
|
||||
|
||||
Discovers available tools from the provided search paths. Each Agent
|
||||
should have its own ToolManager instance.
|
||||
"""
|
||||
|
||||
def __init__(self, config: VibeConfig) -> None:
|
||||
self._config = config
|
||||
self._instances: dict[str, BaseTool] = {}
|
||||
self._search_paths: list[Path] = self._compute_search_paths(config)
|
||||
|
||||
self._available: dict[str, type[BaseTool]] = {
|
||||
cls.get_name(): cls for cls in self._iter_tool_classes(self._search_paths)
|
||||
}
|
||||
self._integrate_mcp()
|
||||
|
||||
@staticmethod
|
||||
def _compute_search_paths(config: VibeConfig) -> list[Path]:
|
||||
paths: list[Path] = [DEFAULT_TOOL_DIR]
|
||||
|
||||
for p in config.tool_paths:
|
||||
path = Path(p).expanduser().resolve()
|
||||
if path.is_dir():
|
||||
paths.append(path)
|
||||
|
||||
cwd = config.effective_workdir
|
||||
for directory in (cwd, *cwd.parents):
|
||||
tools_dir = directory / ".vibe" / "tools"
|
||||
if tools_dir.is_dir():
|
||||
paths.append(tools_dir)
|
||||
break
|
||||
|
||||
global_tools = get_vibe_home() / "tools"
|
||||
if global_tools.is_dir():
|
||||
paths.append(global_tools)
|
||||
|
||||
unique: list[Path] = []
|
||||
seen: set[Path] = set()
|
||||
for p in paths:
|
||||
rp = p.resolve()
|
||||
if rp not in seen:
|
||||
seen.add(rp)
|
||||
unique.append(rp)
|
||||
return unique
|
||||
|
||||
@staticmethod
|
||||
def _iter_tool_classes(search_paths: list[Path]) -> Iterator[type[BaseTool]]:
|
||||
for base in search_paths:
|
||||
if not base.is_dir():
|
||||
continue
|
||||
|
||||
for path in base.rglob("*.py"):
|
||||
if not path.is_file():
|
||||
continue
|
||||
name = path.name
|
||||
if name.startswith("_"):
|
||||
continue
|
||||
|
||||
stem = re.sub(r"[^0-9A-Za-z_]", "_", path.stem) or "mod"
|
||||
module_name = f"vibe_tools_discovered_{stem}"
|
||||
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
if spec is None or spec.loader is None:
|
||||
continue
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
for obj in vars(module).values():
|
||||
if not inspect.isclass(obj):
|
||||
continue
|
||||
if not issubclass(obj, BaseTool) or obj is BaseTool:
|
||||
continue
|
||||
if inspect.isabstract(obj):
|
||||
continue
|
||||
yield obj
|
||||
|
||||
@staticmethod
|
||||
def discover_tool_defaults(
|
||||
search_paths: list[Path] | None = None,
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
if search_paths is None:
|
||||
search_paths = [DEFAULT_TOOL_DIR]
|
||||
|
||||
defaults: dict[str, dict[str, Any]] = {}
|
||||
for cls in ToolManager._iter_tool_classes(search_paths):
|
||||
try:
|
||||
tool_name = cls.get_name()
|
||||
config_class = cls._get_tool_config_class()
|
||||
defaults[tool_name] = config_class().model_dump(exclude_none=True)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to get defaults for tool %s: %s", cls.__name__, e
|
||||
)
|
||||
continue
|
||||
return defaults
|
||||
|
||||
def available_tools(self) -> dict[str, type[BaseTool]]:
|
||||
return dict(self._available)
|
||||
|
||||
def _integrate_mcp(self) -> None:
|
||||
if not self._config.mcp_servers:
|
||||
return
|
||||
run_sync(self._integrate_mcp_async())
|
||||
|
||||
async def _integrate_mcp_async(self) -> None:
|
||||
try:
|
||||
http_count = 0
|
||||
stdio_count = 0
|
||||
|
||||
for srv in self._config.mcp_servers:
|
||||
match srv.transport:
|
||||
case "http" | "streamable-http":
|
||||
http_count += await self._register_http_server(srv)
|
||||
case "stdio":
|
||||
stdio_count += await self._register_stdio_server(srv)
|
||||
case _:
|
||||
logger.warning("Unsupported MCP transport: %r", srv.transport)
|
||||
|
||||
logger.info(
|
||||
"MCP integration registered %d tools (http=%d, stdio=%d)",
|
||||
http_count + stdio_count,
|
||||
http_count,
|
||||
stdio_count,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to integrate MCP tools: %s", exc)
|
||||
|
||||
async def _register_http_server(self, srv: MCPHttp | MCPStreamableHttp) -> int:
|
||||
url = (srv.url or "").strip()
|
||||
if not url:
|
||||
logger.warning("MCP server '%s' missing url for http transport", srv.name)
|
||||
return 0
|
||||
|
||||
headers = srv.http_headers()
|
||||
try:
|
||||
tools: list[RemoteTool] = await list_tools_http(url, headers=headers)
|
||||
except Exception as exc:
|
||||
logger.warning("MCP HTTP discovery failed for %s: %s", url, exc)
|
||||
return 0
|
||||
|
||||
added = 0
|
||||
for remote in tools:
|
||||
try:
|
||||
proxy_cls = create_mcp_http_proxy_tool_class(
|
||||
url=url,
|
||||
remote=remote,
|
||||
alias=srv.name,
|
||||
server_hint=srv.prompt,
|
||||
headers=headers,
|
||||
)
|
||||
self._available[proxy_cls.get_name()] = proxy_cls
|
||||
added += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to register MCP HTTP tool '%s' from %s: %r",
|
||||
getattr(remote, "name", "<unknown>"),
|
||||
url,
|
||||
exc,
|
||||
)
|
||||
return added
|
||||
|
||||
async def _register_stdio_server(self, srv: MCPStdio) -> int:
|
||||
cmd = srv.argv()
|
||||
if not cmd:
|
||||
logger.warning("MCP stdio server '%s' has invalid/empty command", srv.name)
|
||||
return 0
|
||||
|
||||
try:
|
||||
tools: list[RemoteTool] = await list_tools_stdio(cmd)
|
||||
except Exception as exc:
|
||||
logger.warning("MCP stdio discovery failed for %r: %s", cmd, exc)
|
||||
return 0
|
||||
|
||||
added = 0
|
||||
for remote in tools:
|
||||
try:
|
||||
proxy_cls = create_mcp_stdio_proxy_tool_class(
|
||||
command=cmd, remote=remote, alias=srv.name, server_hint=srv.prompt
|
||||
)
|
||||
self._available[proxy_cls.get_name()] = proxy_cls
|
||||
added += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to register MCP stdio tool '%s' from %r: %r",
|
||||
getattr(remote, "name", "<unknown>"),
|
||||
cmd,
|
||||
exc,
|
||||
)
|
||||
return added
|
||||
|
||||
def get_tool_config(self, tool_name: str) -> BaseToolConfig:
|
||||
tool_class = self._available.get(tool_name)
|
||||
|
||||
if tool_class:
|
||||
config_class = tool_class._get_tool_config_class()
|
||||
default_config = config_class()
|
||||
else:
|
||||
config_class = BaseToolConfig
|
||||
default_config = BaseToolConfig()
|
||||
|
||||
user_overrides = self._config.tools.get(tool_name)
|
||||
if user_overrides is None:
|
||||
merged_dict = default_config.model_dump()
|
||||
else:
|
||||
merged_dict = {**default_config.model_dump(), **user_overrides.model_dump()}
|
||||
|
||||
if self._config.workdir is not None:
|
||||
merged_dict["workdir"] = self._config.workdir
|
||||
|
||||
return config_class.model_validate(merged_dict)
|
||||
|
||||
def get(self, tool_name: str) -> BaseTool:
|
||||
"""Get a tool instance, creating it lazily on first call.
|
||||
|
||||
Raises:
|
||||
NoSuchToolError: If the requested tool is not available.
|
||||
"""
|
||||
if tool_name in self._instances:
|
||||
return self._instances[tool_name]
|
||||
|
||||
if tool_name not in self._available:
|
||||
raise NoSuchToolError(
|
||||
f"Unknown tool: {tool_name}. Available: {list(self._available.keys())}"
|
||||
)
|
||||
|
||||
tool_class = self._available[tool_name]
|
||||
tool_config = self.get_tool_config(tool_name)
|
||||
self._instances[tool_name] = tool_class.from_config(tool_config)
|
||||
return self._instances[tool_name]
|
||||
|
||||
def reset_all(self) -> None:
|
||||
self._instances.clear()
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue