v2.15.0 (#773)
Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai> Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com> Co-authored-by: Paul VEZIA <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Guillaume LE GOFF
GitHub
parent
702d0f412e
commit
cafb6d4147
247 changed files with 12401 additions and 3029 deletions
|
|
@ -1,8 +1,27 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum, auto
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from uuid import uuid4
|
||||
|
||||
from textual.widget import Widget
|
||||
|
||||
from vibe.cli.textual_ui.widgets.messages import (
|
||||
BashOutputMessage,
|
||||
ErrorMessage,
|
||||
QueueHeaderMessage,
|
||||
UserMessage,
|
||||
)
|
||||
from vibe.core.autocompletion.path_prompt import PathPromptPayload
|
||||
from vibe.core.logger import logger
|
||||
from vibe.core.types import ImageAttachment
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vibe.core.config import ModelConfig
|
||||
|
||||
|
||||
class QueuedItemKind(StrEnum):
|
||||
|
|
@ -14,16 +33,15 @@ class QueuedItemKind(StrEnum):
|
|||
class QueuedItem:
|
||||
kind: QueuedItemKind
|
||||
content: str
|
||||
skill_name: str | None = None
|
||||
images: list[ImageAttachment] | None = None
|
||||
payload: PathPromptPayload | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class MessageQueue:
|
||||
_items: list[QueuedItem] = field(default_factory=list)
|
||||
_paused: bool = False
|
||||
_on_change: Callable[[], None] | None = None
|
||||
|
||||
def set_change_listener(self, listener: Callable[[], None] | None) -> None:
|
||||
self._on_change = listener
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._items)
|
||||
|
|
@ -39,47 +57,370 @@ class MessageQueue:
|
|||
def paused(self) -> bool:
|
||||
return self._paused
|
||||
|
||||
def append_prompt(self, content: str) -> None:
|
||||
self._items.append(QueuedItem(QueuedItemKind.PROMPT, content))
|
||||
self._notify()
|
||||
def append_prompt(
|
||||
self,
|
||||
content: str,
|
||||
*,
|
||||
skill_name: str | None = None,
|
||||
images: list[ImageAttachment] | None = None,
|
||||
payload: PathPromptPayload | None = None,
|
||||
) -> None:
|
||||
self._items.append(
|
||||
QueuedItem(
|
||||
QueuedItemKind.PROMPT,
|
||||
content,
|
||||
skill_name,
|
||||
images=images,
|
||||
payload=payload,
|
||||
)
|
||||
)
|
||||
|
||||
def append_bash(self, content: str) -> None:
|
||||
self._items.append(QueuedItem(QueuedItemKind.BASH, content))
|
||||
self._notify()
|
||||
|
||||
def prepend_prompts(self, items: list[QueuedItem]) -> None:
|
||||
if not items:
|
||||
return
|
||||
self._items[:0] = items
|
||||
|
||||
def pop_last(self) -> QueuedItem | None:
|
||||
if not self._items:
|
||||
return None
|
||||
item = self._items.pop()
|
||||
self._notify()
|
||||
if not self._items:
|
||||
self._paused = False
|
||||
return item
|
||||
|
||||
def pop_first(self) -> QueuedItem | None:
|
||||
if not self._items:
|
||||
return None
|
||||
item = self._items.pop(0)
|
||||
self._notify()
|
||||
return item
|
||||
return self._items.pop(0)
|
||||
|
||||
def pause(self) -> None:
|
||||
if self._paused:
|
||||
return
|
||||
self._paused = True
|
||||
self._notify()
|
||||
|
||||
def resume(self) -> None:
|
||||
if not self._paused:
|
||||
return
|
||||
self._paused = False
|
||||
self._notify()
|
||||
|
||||
def clear(self) -> None:
|
||||
if not self._items and not self._paused:
|
||||
return
|
||||
self._items.clear()
|
||||
self._paused = False
|
||||
self._notify()
|
||||
|
||||
def _notify(self) -> None:
|
||||
if self._on_change is not None:
|
||||
self._on_change()
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueuePorts:
|
||||
"""Callbacks the controller uses to reach back into the app.
|
||||
|
||||
Everything the drain engine needs that only ``VibeApp`` can provide is
|
||||
funnelled through here, so the controller never touches app internals
|
||||
directly. The app keeps ownership of the things it must (the agent task
|
||||
handle, the loading widget, the remote/feedback managers).
|
||||
"""
|
||||
|
||||
mount_and_scroll: Callable[..., Awaitable[None]]
|
||||
agent_running: Callable[[], bool]
|
||||
bash_task: Callable[[], asyncio.Task | None]
|
||||
active_model: Callable[[], ModelConfig | None]
|
||||
remote_is_active: Callable[[], bool]
|
||||
remote_stop_stream: Callable[[], Awaitable[None]]
|
||||
remove_loading_widget: Callable[[], Awaitable[None]]
|
||||
set_loading_queue_count: Callable[[int], None]
|
||||
inject_user_context: Callable[..., Awaitable[None]]
|
||||
next_message_index: Callable[[], int]
|
||||
start_agent_turn: Callable[..., asyncio.Task]
|
||||
await_agent_turn: Callable[[], Awaitable[None]]
|
||||
run_bash: Callable[..., asyncio.Task]
|
||||
handle_user_message: Callable[[str], Awaitable[None]]
|
||||
maybe_show_feedback_bar: Callable[[], None]
|
||||
send_skill_telemetry: Callable[[str | None], None]
|
||||
send_at_mention_telemetry: Callable[[PathPromptPayload, str], None]
|
||||
render_payload: Callable[[PathPromptPayload], str]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Pending:
|
||||
item: QueuedItem
|
||||
widget: UserMessage
|
||||
|
||||
|
||||
class QueueController:
|
||||
"""Owns the queued-input lifecycle: data, pending widgets, header, drain.
|
||||
|
||||
``MessageQueue`` stays a pure data structure; this controller keeps the
|
||||
parallel list of pending widgets in lockstep with it, manages the header
|
||||
widget, and runs the drain engine that turns queued items into real turns.
|
||||
"""
|
||||
|
||||
def __init__(self, ports: QueuePorts) -> None:
|
||||
self._ports = ports
|
||||
self._queue = MessageQueue()
|
||||
self._widgets: list[Widget] = []
|
||||
self._header: QueueHeaderMessage | None = None
|
||||
self._drain_task: asyncio.Task | None = None
|
||||
|
||||
@property
|
||||
def queue(self) -> MessageQueue:
|
||||
return self._queue
|
||||
|
||||
@property
|
||||
def header(self) -> QueueHeaderMessage | None:
|
||||
return self._header
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self._queue)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._queue)
|
||||
|
||||
# -- pin target (used by the app's _mount_and_scroll) ------------------
|
||||
|
||||
def pin_target(self, messages_area: Widget) -> Widget | None:
|
||||
target: Widget | None = self._header
|
||||
if target is None and self._widgets:
|
||||
target = self._widgets[0]
|
||||
if target is not None and target.parent is messages_area:
|
||||
return target
|
||||
return None
|
||||
|
||||
def _last_queue_anchor(self) -> Widget | None:
|
||||
if self._widgets:
|
||||
return self._widgets[-1]
|
||||
return self._header
|
||||
|
||||
# -- quit / count helpers --------------------------------------------
|
||||
|
||||
def quit_warning_extra(self) -> str:
|
||||
if not self._queue:
|
||||
return ""
|
||||
n = len(self._queue)
|
||||
plural = "s" if n != 1 else ""
|
||||
return f"{n} queued message{plural} will be discarded"
|
||||
|
||||
def _push_loading_queue_count(self) -> None:
|
||||
self._ports.set_loading_queue_count(len(self._queue))
|
||||
|
||||
def notify_busy_changed(self) -> None:
|
||||
self._push_loading_queue_count()
|
||||
|
||||
# -- enqueue ----------------------------------------------------------
|
||||
|
||||
async def enqueue_prompt(
|
||||
self,
|
||||
content: str,
|
||||
*,
|
||||
skill_name: str | None = None,
|
||||
images: list[ImageAttachment] | None = None,
|
||||
payload: PathPromptPayload | None = None,
|
||||
) -> None:
|
||||
self._queue.append_prompt(
|
||||
content, skill_name=skill_name, images=images, payload=payload
|
||||
)
|
||||
await self._ensure_header()
|
||||
widget = UserMessage(content, pending=True, images=images or None)
|
||||
anchor = self._last_queue_anchor()
|
||||
self._widgets.append(widget)
|
||||
await self._ports.mount_and_scroll(widget, after=anchor)
|
||||
self._push_loading_queue_count()
|
||||
|
||||
async def enqueue_bash(self, content: str) -> None:
|
||||
self._queue.append_bash(content)
|
||||
await self._ensure_header()
|
||||
widget = BashOutputMessage(content, str(Path.cwd()), pending=True)
|
||||
widget.set_queued(True)
|
||||
anchor = self._last_queue_anchor()
|
||||
self._widgets.append(widget)
|
||||
await self._ports.mount_and_scroll(widget, after=anchor)
|
||||
self._push_loading_queue_count()
|
||||
|
||||
async def pop_last(self) -> bool:
|
||||
item = self._queue.pop_last()
|
||||
if item is None:
|
||||
return False
|
||||
widget = self._widgets.pop() if self._widgets else None
|
||||
if widget is not None:
|
||||
await widget.remove()
|
||||
await self._remove_header_if_empty()
|
||||
self._push_loading_queue_count()
|
||||
return True
|
||||
|
||||
# -- header lifecycle -------------------------------------------------
|
||||
|
||||
async def _ensure_header(self) -> None:
|
||||
if self._header is not None:
|
||||
return
|
||||
header = QueueHeaderMessage(paused=self._queue.paused)
|
||||
self._header = header
|
||||
await self._ports.mount_and_scroll(header)
|
||||
|
||||
async def _remove_header_if_empty(self) -> None:
|
||||
if self._queue or self._header is None:
|
||||
return
|
||||
await self._remove_header()
|
||||
|
||||
async def _remove_header(self) -> None:
|
||||
if self._header is None:
|
||||
return
|
||||
header = self._header
|
||||
self._header = None
|
||||
await header.remove()
|
||||
|
||||
def set_paused(self, paused: bool) -> None:
|
||||
if paused:
|
||||
self._queue.pause()
|
||||
else:
|
||||
self._queue.resume()
|
||||
if self._header is not None:
|
||||
self._header.set_paused(self._queue.paused)
|
||||
|
||||
# -- drain engine -----------------------------------------------------
|
||||
|
||||
def start_drain_if_needed(self) -> None:
|
||||
if self._drain_task is not None and not self._drain_task.done():
|
||||
return
|
||||
if not self._queue or self._queue.paused:
|
||||
return
|
||||
if self._ports.agent_running():
|
||||
return
|
||||
bash_task = self._ports.bash_task()
|
||||
if bash_task is not None and not bash_task.done():
|
||||
return
|
||||
self._drain_task = asyncio.create_task(self._drain())
|
||||
|
||||
@property
|
||||
def draining(self) -> bool:
|
||||
return self._drain_task is not None and not self._drain_task.done()
|
||||
|
||||
async def _drain(self) -> None:
|
||||
try:
|
||||
while self._queue and not self._queue.paused:
|
||||
await self._remove_header()
|
||||
pending = await self._consume_until_bash_or_empty()
|
||||
if not pending:
|
||||
continue
|
||||
if self._queue.paused:
|
||||
self._requeue(pending)
|
||||
continue
|
||||
await self._run_pending_as_llm_turn(pending)
|
||||
except Exception:
|
||||
logger.exception("Queue drain crashed")
|
||||
finally:
|
||||
self._drain_task = None
|
||||
self.notify_busy_changed()
|
||||
await self._remove_header_if_empty()
|
||||
|
||||
async def _consume_until_bash_or_empty(self) -> list[_Pending]:
|
||||
pending: list[_Pending] = []
|
||||
while self._queue and not self._queue.paused:
|
||||
item = self._queue.pop_first()
|
||||
if item is None:
|
||||
break
|
||||
widget = self._widgets.pop(0) if self._widgets else None
|
||||
if item.kind == QueuedItemKind.BASH:
|
||||
await self._flush_pending_prompts(pending)
|
||||
pending = []
|
||||
bash_widget = widget if isinstance(widget, BashOutputMessage) else None
|
||||
if not await self._run_bash(item.content, bash_widget):
|
||||
return []
|
||||
elif isinstance(widget, UserMessage):
|
||||
pending.append(_Pending(item, widget))
|
||||
return pending
|
||||
|
||||
def _requeue(self, pending: list[_Pending]) -> None:
|
||||
self._queue.prepend_prompts([p.item for p in pending])
|
||||
self._widgets[:0] = [p.widget for p in pending]
|
||||
|
||||
async def _run_pending_as_llm_turn(self, pending: list[_Pending]) -> None:
|
||||
if not await self._gate_queued_images_for_vision(pending):
|
||||
return
|
||||
head, tail = pending[:-1], pending[-1]
|
||||
for p in head:
|
||||
await self._inject_head_item(p.item, p.widget)
|
||||
await p.widget.set_pending(False)
|
||||
self._link_consecutive_user_messages([p.widget for p in pending])
|
||||
await self._run_tail_prompt(tail.item, tail.widget)
|
||||
await self._await_tail_turn()
|
||||
|
||||
async def _await_tail_turn(self) -> None:
|
||||
try:
|
||||
await self._ports.await_agent_turn()
|
||||
except asyncio.CancelledError:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
self._push_loading_queue_count()
|
||||
|
||||
async def _flush_pending_prompts(self, pending: list[_Pending]) -> None:
|
||||
if not await self._gate_queued_images_for_vision(pending):
|
||||
return
|
||||
for p in pending:
|
||||
await self._inject_head_item(p.item, p.widget)
|
||||
await p.widget.set_pending(False)
|
||||
self._link_consecutive_user_messages([p.widget for p in pending])
|
||||
|
||||
async def _gate_queued_images_for_vision(self, pending: list[_Pending]) -> bool:
|
||||
if not any(p.item.images for p in pending):
|
||||
return True
|
||||
active_model = self._ports.active_model()
|
||||
if active_model is None or active_model.supports_images:
|
||||
return True
|
||||
self._requeue(pending)
|
||||
self.set_paused(True)
|
||||
await self._ensure_header()
|
||||
await self._ports.mount_and_scroll(
|
||||
ErrorMessage(
|
||||
f"Model `{active_model.alias}` does not support images. "
|
||||
f"Switch with /model, then press Enter to resume the queue.",
|
||||
show_border=False,
|
||||
)
|
||||
)
|
||||
return False
|
||||
|
||||
async def _inject_head_item(self, item: QueuedItem, widget: UserMessage) -> None:
|
||||
widget.message_index = self._ports.next_message_index()
|
||||
message_id = str(uuid4()) if item.payload is not None else None
|
||||
if item.payload is not None:
|
||||
rendered = self._ports.render_payload(item.payload)
|
||||
else:
|
||||
rendered = item.content
|
||||
await self._ports.inject_user_context(
|
||||
rendered, as_message=True, images=item.images, client_message_id=message_id
|
||||
)
|
||||
self._ports.send_skill_telemetry(item.skill_name)
|
||||
if item.payload is not None and message_id is not None:
|
||||
self._ports.send_at_mention_telemetry(item.payload, message_id)
|
||||
|
||||
async def _run_tail_prompt(self, item: QueuedItem, widget: UserMessage) -> None:
|
||||
if self._ports.remote_is_active():
|
||||
await widget.remove()
|
||||
await self._ports.handle_user_message(item.content)
|
||||
self._ports.send_skill_telemetry(item.skill_name)
|
||||
return
|
||||
|
||||
widget.message_index = self._ports.next_message_index()
|
||||
await widget.set_pending(False)
|
||||
self._ports.maybe_show_feedback_bar()
|
||||
|
||||
await self._ports.remote_stop_stream()
|
||||
await self._ports.remove_loading_widget()
|
||||
self._ports.start_agent_turn(
|
||||
item.content, prebuilt_images=item.images, prebuilt_payload=item.payload
|
||||
)
|
||||
self._ports.send_skill_telemetry(item.skill_name)
|
||||
self.notify_busy_changed()
|
||||
|
||||
async def _run_bash(self, command: str, widget: BashOutputMessage | None) -> bool:
|
||||
if widget is not None:
|
||||
widget.set_queued(False)
|
||||
bash_task = self._ports.run_bash(command, existing_widget=widget)
|
||||
self.notify_busy_changed()
|
||||
try:
|
||||
await bash_task
|
||||
except asyncio.CancelledError:
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _link_consecutive_user_messages(widgets: list[UserMessage]) -> None:
|
||||
for prev, curr in zip(widgets, widgets[1:], strict=False):
|
||||
prev.set_show_separator(False)
|
||||
curr.set_follows_previous(True)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue