265 lines
7.8 KiB
Python
265 lines
7.8 KiB
Python
import re
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Literal
|
||
|
||
from pydantic_ai.messages import BinaryContent
|
||
from rich.style import Style
|
||
from textual import on
|
||
from textual.app import ComposeResult
|
||
from textual.binding import Binding
|
||
from textual.containers import Horizontal
|
||
from textual.css.query import NoMatches
|
||
from textual.message import Message
|
||
from textual.reactive import reactive
|
||
from textual.widget import Widget
|
||
from textual.widgets import Static, TextArea
|
||
|
||
from haiku.rag.chat.widgets.image_select import ImageAdded, ImageSelect
|
||
from haiku.rag.utils import image_binary_content
|
||
|
||
MAX_PROMPT_LINES = 10
|
||
|
||
IMAGE_TOKEN_RE = re.compile(r"\[Image #(\d+)\]")
|
||
_IMAGE_TOKEN_HIGHLIGHT = "image-token"
|
||
_IMAGE_TOKEN_STYLE = Style(color="bright_cyan", bold=True)
|
||
|
||
|
||
def build_user_prompt(
|
||
text: str, images: list[bytes]
|
||
) -> str | list[str | BinaryContent]:
|
||
"""Interleave text and images by ``[Image #N]`` tokens, 1-indexed.
|
||
|
||
Without tokens, images are appended after the text. Out-of-range tokens
|
||
stay as literal text.
|
||
"""
|
||
matches = list(IMAGE_TOKEN_RE.finditer(text))
|
||
if not matches:
|
||
if not images:
|
||
return text
|
||
parts: list[str | BinaryContent] = [text] if text else []
|
||
parts.extend(image_binary_content(data) for data in images)
|
||
return parts
|
||
|
||
parts = []
|
||
last = 0
|
||
for m in matches:
|
||
if m.start() > last:
|
||
parts.append(text[last : m.start()])
|
||
idx = int(m.group(1))
|
||
if 1 <= idx <= len(images):
|
||
parts.append(image_binary_content(images[idx - 1]))
|
||
else:
|
||
parts.append(m.group(0))
|
||
last = m.end()
|
||
if last < len(text):
|
||
parts.append(text[last:])
|
||
if not any(isinstance(p, BinaryContent) for p in parts):
|
||
return text
|
||
return parts
|
||
|
||
|
||
class PostableTextArea(TextArea):
|
||
"""TextArea that auto-grows with content, submits on Enter, newline on Shift+Enter."""
|
||
|
||
BINDINGS = TextArea.BINDINGS + [
|
||
Binding(
|
||
key="enter",
|
||
action="submit",
|
||
description="submit",
|
||
show=True,
|
||
key_display=None,
|
||
priority=True,
|
||
),
|
||
Binding(
|
||
key="shift+enter",
|
||
action="newline",
|
||
description="newline",
|
||
show=True,
|
||
key_display=None,
|
||
priority=True,
|
||
id="newline",
|
||
),
|
||
Binding(
|
||
key="ctrl+m",
|
||
action="newline",
|
||
description="newline",
|
||
show=False,
|
||
key_display=None,
|
||
priority=True,
|
||
),
|
||
]
|
||
|
||
@dataclass
|
||
class Submitted(Message):
|
||
input: "PostableTextArea"
|
||
value: str
|
||
|
||
@property
|
||
def control(self) -> "PostableTextArea":
|
||
return self.input
|
||
|
||
def on_mount(self) -> None:
|
||
self.soft_wrap = True
|
||
self._resize_to_content()
|
||
if self._theme is not None: # pragma: no branch
|
||
self._theme.syntax_styles[_IMAGE_TOKEN_HIGHLIGHT] = _IMAGE_TOKEN_STYLE
|
||
self._build_highlight_map()
|
||
self.refresh()
|
||
|
||
def _resize_to_content(self) -> None:
|
||
line_count = max(self.wrapped_document.height, 1)
|
||
self.styles.height = min(line_count, MAX_PROMPT_LINES)
|
||
|
||
def _build_highlight_map(self) -> None:
|
||
super()._build_highlight_map()
|
||
for line_idx in range(self.document.line_count):
|
||
line = self.document.get_line(line_idx)
|
||
for m in IMAGE_TOKEN_RE.finditer(line):
|
||
self._highlights[line_idx].append(
|
||
(m.start(), m.end(), _IMAGE_TOKEN_HIGHLIGHT)
|
||
)
|
||
|
||
def action_submit(self) -> None:
|
||
self.post_message(PostableTextArea.Submitted(self, self.text))
|
||
|
||
def action_newline(self) -> None:
|
||
self.insert("\n")
|
||
|
||
def action_delete_left(self) -> None:
|
||
if self.selection.start != self.selection.end:
|
||
super().action_delete_left()
|
||
return
|
||
span = self._image_token_span_at_cursor("left")
|
||
if span is not None:
|
||
self.delete(*span)
|
||
return
|
||
super().action_delete_left()
|
||
|
||
def action_delete_right(self) -> None:
|
||
if self.selection.start != self.selection.end:
|
||
super().action_delete_right()
|
||
return
|
||
span = self._image_token_span_at_cursor("right")
|
||
if span is not None:
|
||
self.delete(*span)
|
||
return
|
||
super().action_delete_right()
|
||
|
||
def _image_token_span_at_cursor(
|
||
self, direction: Literal["left", "right"]
|
||
) -> tuple[tuple[int, int], tuple[int, int]] | None:
|
||
row, col = self.cursor_location
|
||
line = self.document.get_line(row)
|
||
for m in IMAGE_TOKEN_RE.finditer(line):
|
||
s, e = m.start(), m.end()
|
||
if direction == "left" and s < col <= e:
|
||
return (row, s), (row, e)
|
||
if direction == "right" and s <= col < e:
|
||
return (row, s), (row, e)
|
||
return None
|
||
|
||
|
||
class FlexibleInput(Widget):
|
||
"""Prompt input with image attachment via ctrl+i."""
|
||
|
||
text = reactive("")
|
||
|
||
BINDINGS = [
|
||
Binding("ctrl+i", "add_image", "add image", id="add.image"),
|
||
]
|
||
|
||
DEFAULT_CSS = """
|
||
FlexibleInput {
|
||
height: auto;
|
||
padding: 0 1 1 1;
|
||
border-top: solid $primary-darken-1;
|
||
}
|
||
|
||
FlexibleInput:focus-within {
|
||
border-top: solid $primary;
|
||
}
|
||
|
||
FlexibleInput > Horizontal {
|
||
height: auto;
|
||
}
|
||
|
||
FlexibleInput #promptMarker {
|
||
width: 2;
|
||
height: 1;
|
||
color: $primary;
|
||
}
|
||
|
||
FlexibleInput #promptArea {
|
||
background: transparent;
|
||
border: none;
|
||
padding: 0;
|
||
}
|
||
|
||
FlexibleInput #promptArea > .text-area--cursor-line {
|
||
background: transparent;
|
||
}
|
||
"""
|
||
|
||
@dataclass
|
||
class Submitted(Message):
|
||
input: "FlexibleInput"
|
||
value: str
|
||
|
||
@property
|
||
def control(self) -> "FlexibleInput":
|
||
return self.input
|
||
|
||
def __init__(self, text: str = "", *args, **kwargs) -> None:
|
||
super().__init__(*args, **kwargs)
|
||
self.text = text
|
||
|
||
def on_mount(self) -> None:
|
||
textarea = self.query_one("#promptArea", PostableTextArea)
|
||
textarea.show_line_numbers = False
|
||
textarea.focus()
|
||
|
||
def clear(self) -> None:
|
||
self.text = ""
|
||
self.query_one("#promptArea", PostableTextArea).text = ""
|
||
|
||
def focus(self, scroll_visible: bool = True) -> "FlexibleInput":
|
||
self.query_one("#promptArea", PostableTextArea).focus()
|
||
return self
|
||
|
||
def insert_at_cursor(self, text: str) -> None:
|
||
self.query_one("#promptArea", PostableTextArea).insert(text)
|
||
|
||
def watch_text(self) -> None:
|
||
try:
|
||
textarea = self.query_one("#promptArea", PostableTextArea)
|
||
if textarea.text != self.text:
|
||
textarea.text = self.text
|
||
except NoMatches:
|
||
pass
|
||
|
||
def action_add_image(self) -> None:
|
||
async def on_image_selected(image: tuple[Path, bytes] | None) -> None:
|
||
if image is None:
|
||
return
|
||
path, data = image
|
||
self.post_message(ImageAdded(path, data))
|
||
|
||
self.app.push_screen(ImageSelect(), on_image_selected)
|
||
|
||
@on(PostableTextArea.Submitted, "#promptArea")
|
||
def on_textarea_submitted(self, event: PostableTextArea.Submitted) -> None:
|
||
self.post_message(self.Submitted(self, event.input.text))
|
||
event.stop()
|
||
event.prevent_default()
|
||
|
||
@on(TextArea.Changed, "#promptArea")
|
||
def on_area_changed(self, event: TextArea.Changed) -> None:
|
||
self.text = event.text_area.text
|
||
if isinstance(event.text_area, PostableTextArea): # pragma: no branch
|
||
event.text_area._resize_to_content()
|
||
|
||
def compose(self) -> ComposeResult:
|
||
with Horizontal():
|
||
yield Static("❯", id="promptMarker")
|
||
yield PostableTextArea(id="promptArea")
|