3b8f65b306
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Lucas Marandat <31749711+lucasmrdt@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
import json
|
|
import sys
|
|
from typing import TextIO
|
|
|
|
from vibe.core.teleport.types import (
|
|
TeleportAuthCompleteEvent,
|
|
TeleportAuthRequiredEvent,
|
|
TeleportCheckingGitEvent,
|
|
TeleportCompleteEvent,
|
|
TeleportFetchingUrlEvent,
|
|
TeleportPushingEvent,
|
|
TeleportPushRequiredEvent,
|
|
TeleportStartingWorkflowEvent,
|
|
TeleportWaitingForGitHubEvent,
|
|
)
|
|
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 _print(self, text: str) -> None:
|
|
print(text, file=self.stream)
|
|
|
|
def on_event(self, event: BaseEvent) -> None:
|
|
match event:
|
|
case AssistantEvent():
|
|
self._final_response = event.content
|
|
case TeleportCheckingGitEvent():
|
|
self._print("Preparing workspace...")
|
|
case TeleportPushRequiredEvent(unpushed_count=count):
|
|
self._print(f"Pushing {count} commit(s)...")
|
|
case TeleportPushingEvent():
|
|
self._print("Syncing with remote...")
|
|
case TeleportStartingWorkflowEvent():
|
|
self._print("Teleporting...")
|
|
case TeleportWaitingForGitHubEvent(message=msg):
|
|
self._print(msg or "Connecting to GitHub...")
|
|
case TeleportAuthRequiredEvent(oauth_url=url, message=msg):
|
|
self._print(msg or f"Open to authorize GitHub: {url}")
|
|
case TeleportAuthCompleteEvent():
|
|
self._print("GitHub authorized")
|
|
case TeleportFetchingUrlEvent():
|
|
self._print("Finalizing...")
|
|
case TeleportCompleteEvent():
|
|
self._final_response = event.url
|
|
|
|
def finalize(self) -> str | None:
|
|
return self._final_response
|
|
|
|
|
|
class JsonOutputFormatter(OutputFormatter):
|
|
def on_message_added(self, message: LLMMessage) -> None:
|
|
self._messages.append(message)
|
|
|
|
def on_event(self, event: BaseEvent) -> None:
|
|
pass
|
|
|
|
def finalize(self) -> str | None:
|
|
messages_data = [msg.model_dump(mode="json") for msg in self._messages]
|
|
json.dump(messages_data, self.stream, indent=2, ensure_ascii=False)
|
|
self.stream.write("\n")
|
|
self.stream.flush()
|
|
return None
|
|
|
|
|
|
class StreamingJsonOutputFormatter(OutputFormatter):
|
|
def on_message_added(self, message: LLMMessage) -> None:
|
|
json.dump(message.model_dump(mode="json"), self.stream, ensure_ascii=False)
|
|
self.stream.write("\n")
|
|
self.stream.flush()
|
|
|
|
def on_event(self, event: BaseEvent) -> None:
|
|
pass
|
|
|
|
def finalize(self) -> str | None:
|
|
return None
|
|
|
|
|
|
def create_formatter(
|
|
format_type: OutputFormat, stream: TextIO = sys.stdout
|
|
) -> OutputFormatter:
|
|
formatters = {
|
|
OutputFormat.TEXT: TextOutputFormatter,
|
|
OutputFormat.JSON: JsonOutputFormatter,
|
|
OutputFormat.STREAMING: StreamingJsonOutputFormatter,
|
|
}
|
|
|
|
formatter_class = formatters.get(format_type, TextOutputFormatter)
|
|
return formatter_class(stream)
|