diff --git a/AgentCrew/modules/console/conversation_browser/browser_ui.py b/AgentCrew/modules/console/conversation_browser/browser_ui.py index 3f05e19b..32ee4fde 100644 --- a/AgentCrew/modules/console/conversation_browser/browser_ui.py +++ b/AgentCrew/modules/console/conversation_browser/browser_ui.py @@ -27,6 +27,7 @@ RICH_STYLE_YELLOW, RICH_STYLE_YELLOW_BOLD, ) +from .search import ConversationSearchIndex, SearchMatch class ConversationBrowserUI: @@ -44,6 +45,7 @@ def __init__( self.scroll_offset = 0 self._get_conversation_history = get_conversation_history self._preview_cache: dict[str, tuple[list[dict[str, Any]], int]] = {} + self._search_index = ConversationSearchIndex(get_conversation_history) self.selected_items: set[int] = set() self._live: Live | None = None self._layout: Layout | None = None @@ -56,11 +58,12 @@ def max_list_items(self) -> int: def set_conversations(self, conversations: list[dict[str, Any]]): """Set the conversations list to browse.""" - self._all_conversations = conversations - self.conversations = conversations + self._all_conversations = list(conversations) + self.conversations = list(conversations) self.selected_index = 0 self.scroll_offset = 0 self._preview_cache.clear() + self._search_index.clear() self.selected_items.clear() self._search_query = "" self._search_mode = False @@ -82,7 +85,7 @@ def exit_search_mode(self, clear_filter: bool = False): self._search_mode = False if clear_filter: self._search_query = "" - self.conversations = self._all_conversations + self.conversations = list(self._all_conversations) self.selected_index = 0 self.scroll_offset = 0 self.selected_items.clear() @@ -105,15 +108,10 @@ def backspace_search(self): def _filter_conversations(self): """Filter conversations based on search query.""" - if not self._search_query: - self.conversations = self._all_conversations - else: - query_lower = self._search_query.lower() - self.conversations = [ - c - for c in self._all_conversations - if query_lower in c.get("title", "").lower() - ] + self.conversations = self._search_index.filter( + self._all_conversations, + self._search_query, + ) self.selected_index = 0 self.scroll_offset = 0 self.selected_items.clear() @@ -145,12 +143,27 @@ def get_selected_conversation_ids(self) -> list[str]: def remove_conversations(self, indices: list[int]): """Remove conversations at specified indices and update UI state.""" - for idx in sorted(indices, reverse=True): - if 0 <= idx < len(self.conversations): - convo_id = self.conversations[idx].get("id") - if convo_id: - self._preview_cache.pop(convo_id, None) - del self.conversations[idx] + conversation_ids = { + self.conversations[idx].get("id") + for idx in indices + if 0 <= idx < len(self.conversations) and self.conversations[idx].get("id") + } + if not conversation_ids: + return + + self._all_conversations = [ + conversation + for conversation in self._all_conversations + if conversation.get("id") not in conversation_ids + ] + self.conversations = [ + conversation + for conversation in self.conversations + if conversation.get("id") not in conversation_ids + ] + for conversation_id in conversation_ids: + self._preview_cache.pop(conversation_id, None) + self._search_index.remove(conversation_ids) self.selected_items.clear() if self.selected_index >= len(self.conversations): self.selected_index = max(0, len(self.conversations) - 1) @@ -430,6 +443,13 @@ def _create_preview_panel(self, panel_height: int | None = None) -> Panel: preview_lines.append(Text("")) preview_lines.append(meta_table) + + search_match = ( + self._search_index.get_match(convo_id) if self._search_query else None + ) + if search_match is not None: + preview_lines.extend(self._create_search_match_preview(search_match)) + preview_lines.append(Text("")) preview_lines.append(Rule(title="Recent Messages", style=RICH_STYLE_GRAY)) @@ -491,6 +511,31 @@ def _create_preview_panel(self, panel_height: int | None = None) -> Panel: box=ROUNDED, ) + def _create_search_match_preview(self, match: SearchMatch) -> list[Text | Rule]: + """Render the first title or message match with surrounding context.""" + source_label = "Title" + if match.source == "message": + source_label = (match.role or "message").capitalize() + + header = Text() + header.append(f"{source_label}: ", style=RICH_STYLE_GREEN_BOLD) + + snippet = match.snippet() + if snippet.has_leading_ellipsis: + header.append("…", style=RICH_STYLE_GRAY) + header.append(snippet.before, style=RICH_STYLE_WHITE) + header.append(snippet.matched, style="bold black on yellow") + header.append(snippet.after, style=RICH_STYLE_WHITE) + if snippet.has_trailing_ellipsis: + header.append("…", style=RICH_STYLE_GRAY) + + return [ + Text(""), + Rule(title="Search Match", style=RICH_STYLE_YELLOW), + Text(""), + header, + ] + def _create_help_panel(self) -> Panel: """Create the help panel with keyboard shortcuts.""" if self._search_mode: @@ -569,7 +614,7 @@ def _create_search_bar(self) -> Panel: search_table, border_style="cyan", box=ROUNDED, - title=Text("Search ", style=RICH_STYLE_YELLOW_BOLD), + title=Text("Search Titles & Messages ", style=RICH_STYLE_YELLOW_BOLD), ) def _create_layout(self) -> Layout: diff --git a/AgentCrew/modules/console/conversation_browser/search.py b/AgentCrew/modules/console/conversation_browser/search.py new file mode 100644 index 00000000..65c01748 --- /dev/null +++ b/AgentCrew/modules/console/conversation_browser/search.py @@ -0,0 +1,278 @@ +"""Full-text search support for the console conversation browser.""" + +from __future__ import annotations + +import re +from array import array +from bisect import bisect_right +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from typing import Any + +from loguru import logger + +_INTERNAL_CONTENT_PREFIXES = ( + "Memories related to the user request:", + "Content of ", +) + + +@dataclass(frozen=True) +class SearchFragment: + """A user-visible piece of searchable conversation text.""" + + role: str + text: str + normalized_text: str = field(init=False, repr=False, compare=False) + _extra_normalized_offsets: array[int] | None = field( + init=False, + default=None, + repr=False, + compare=False, + ) + + def __post_init__(self) -> None: + normalized_text = self.text.casefold() + object.__setattr__(self, "normalized_text", normalized_text) + + if len(normalized_text) == len(self.text): + return + + extra_normalized_offsets = array("I") + normalized_offset = 0 + for char in self.text: + folded_length = len(char.casefold()) + extra_normalized_offsets.extend( + range(normalized_offset + 1, normalized_offset + folded_length) + ) + normalized_offset += folded_length + + object.__setattr__( + self, + "_extra_normalized_offsets", + extra_normalized_offsets, + ) + + def find_span(self, normalized_query: str) -> tuple[int, int] | None: + """Find a normalized query and map its span back to the original text.""" + normalized_start = self.normalized_text.find(normalized_query) + if normalized_start < 0: + return None + + normalized_end = normalized_start + len(normalized_query) + if self._extra_normalized_offsets is None: + return normalized_start, normalized_end + + original_start = normalized_start - bisect_right( + self._extra_normalized_offsets, + normalized_start, + ) + last_normalized_offset = normalized_end - 1 + original_end = ( + last_normalized_offset + - bisect_right( + self._extra_normalized_offsets, + last_normalized_offset, + ) + + 1 + ) + return original_start, original_end + + +@dataclass(frozen=True) +class SearchSnippet: + """Display-ready context surrounding a search match.""" + + before: str + matched: str + after: str + has_leading_ellipsis: bool + has_trailing_ellipsis: bool + + +@dataclass(frozen=True) +class SearchMatch: + """The first title or message match for a conversation.""" + + source: str + text: str + start: int + end: int + role: str | None = None + + def snippet(self, context_chars: int = 60) -> SearchSnippet: + """Return compact display text around the matched range.""" + context_start = max(0, self.start - context_chars) + context_end = min(len(self.text), self.end + context_chars) + + before = _collapse_whitespace(self.text[context_start : self.start]) + matched = _collapse_whitespace(self.text[self.start : self.end]) + after = _collapse_whitespace(self.text[self.end : context_end]) + + return SearchSnippet( + before=before, + matched=matched, + after=after, + has_leading_ellipsis=context_start > 0, + has_trailing_ellipsis=context_end < len(self.text), + ) + + +def _is_searchable_text(text: str) -> bool: + stripped = text.strip() + return bool(stripped) and not stripped.startswith(_INTERNAL_CONTENT_PREFIXES) + + +def _collapse_whitespace(text: str) -> str: + return re.sub(r"\s+", " ", text) + + +def _find_casefold_span(text: str, normalized_query: str) -> tuple[int, int] | None: + """Map a casefolded match back to offsets in the original text.""" + normalized_chars: list[str] = [] + original_offsets: list[int] = [] + for original_offset, char in enumerate(text): + folded = char.casefold() + normalized_chars.append(folded) + original_offsets.extend([original_offset] * len(folded)) + + normalized_text = "".join(normalized_chars) + normalized_start = normalized_text.find(normalized_query) + if normalized_start < 0: + return None + + normalized_end = normalized_start + len(normalized_query) + original_start = original_offsets[normalized_start] + original_end = original_offsets[normalized_end - 1] + 1 + return original_start, original_end + + +def iter_searchable_message_fragments(message: Any) -> Iterable[SearchFragment]: + """Yield searchable fragments with their originating chat role.""" + if not isinstance(message, dict) or message.get("role") not in { + "user", + "assistant", + }: + return + + role = message["role"] + content = message.get("content", "") + if isinstance(content, str): + if _is_searchable_text(content): + yield SearchFragment(role=role, text=content) + return + + if not isinstance(content, list): + return + + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = block.get("text", "") + if isinstance(text, str) and _is_searchable_text(text): + yield SearchFragment(role=role, text=text) + + +class ConversationSearchIndex: + """ + Session-local full-text index for persisted conversations. + + Conversation histories are loaded lazily and normalized once. Subsequent queries search the in-memory index without repeating filesystem reads. + """ + + def __init__( + self, + get_conversation_history: ( + Callable[[str], list[dict[str, Any]] | None] | None + ) = None, + ) -> None: + self._get_conversation_history = get_conversation_history + self._message_fragments: dict[str, list[SearchFragment]] = {} + self._matches: dict[str, SearchMatch] = {} + + def clear(self) -> None: + """Discard all indexed conversation content.""" + self._message_fragments.clear() + self._matches.clear() + + def remove(self, conversation_ids: Iterable[str]) -> None: + """Remove conversations from the in-memory index.""" + for conversation_id in conversation_ids: + self._message_fragments.pop(conversation_id, None) + self._matches.pop(conversation_id, None) + + def get_match(self, conversation_id: str) -> SearchMatch | None: + """Return the match produced by the most recent filter operation.""" + return self._matches.get(conversation_id) + + def filter( + self, + conversations: list[dict[str, Any]], + query: str, + ) -> list[dict[str, Any]]: + """ + Return conversations whose title/visible messages match the user query. + """ + normalized_query = query.casefold() + self._matches.clear() + if not normalized_query: + return list(conversations) + + matches: list[dict[str, Any]] = [] + for conversation in conversations: + conversation_id = conversation.get("id") + if not isinstance(conversation_id, str) or not conversation_id: + continue + + title = conversation.get("title", "") + title_span = ( + _find_casefold_span(title, normalized_query) + if isinstance(title, str) + else None + ) + if title_span is not None: + matches.append(conversation) + self._matches[conversation_id] = SearchMatch( + source="title", + text=title, + start=title_span[0], + end=title_span[1], + ) + continue + + for fragment in self._get_message_fragments(conversation_id): + message_span = fragment.find_span(normalized_query) + if message_span is None: + continue + matches.append(conversation) + self._matches[conversation_id] = SearchMatch( + source="message", + role=fragment.role, + text=fragment.text, + start=message_span[0], + end=message_span[1], + ) + break + + return matches + + def _get_message_fragments(self, conversation_id: str) -> list[SearchFragment]: + if conversation_id in self._message_fragments: + return self._message_fragments[conversation_id] + + fragments: list[SearchFragment] = [] + if self._get_conversation_history is not None: + try: + history = self._get_conversation_history(conversation_id) + if isinstance(history, list): + for message in history: + fragments.extend(iter_searchable_message_fragments(message)) + except Exception as exc: + logger.warning( + "Error indexing conversation '{}' for search: {}", + conversation_id, + exc, + ) + + self._message_fragments[conversation_id] = fragments + return fragments diff --git a/tests/console/test_conversation_search.py b/tests/console/test_conversation_search.py new file mode 100644 index 00000000..10c6b515 --- /dev/null +++ b/tests/console/test_conversation_search.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +from typing import Any + +from rich.console import Console + +from AgentCrew.modules.console.conversation_browser.browser_ui import ( + ConversationBrowserUI, +) +from AgentCrew.modules.console.conversation_browser.search import ( + ConversationSearchIndex, + SearchFragment, + SearchMatch, + iter_searchable_message_fragments, +) + + +def _conversation(conversation_id: str, title: str) -> dict[str, Any]: + return {"id": conversation_id, "title": title, "timestamp": "2026-01-01"} + + +def test_filters_by_title_without_loading_matching_history() -> None: + calls: list[str] = [] + + def load_history(conversation_id: str) -> list[dict[str, Any]]: + calls.append(conversation_id) + return [] + + index = ConversationSearchIndex(load_history) + conversations = [_conversation("one", "Database design")] + + assert index.filter(conversations, "DATABASE") == conversations + assert calls == [] + assert index.get_match("one") == SearchMatch( + source="title", + text="Database design", + start=0, + end=8, + ) + + +def test_filters_user_and_assistant_message_text() -> None: + histories = { + "one": [{"role": "user", "content": "Plan a release pipeline"}], + "two": [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "First block"}, + {"type": "tool_use", "name": "ignored"}, + {"type": "text", "text": "Kubernetes deployment"}, + ], + } + ], + } + index = ConversationSearchIndex(histories.get) + conversations = [ + _conversation("one", "Release"), + _conversation("two", "Infrastructure"), + ] + + assert [c["id"] for c in index.filter(conversations, "pipeline")] == ["one"] + assert [c["id"] for c in index.filter(conversations, "kubernetes")] == ["two"] + assert index.get_match("two") == SearchMatch( + source="message", + role="assistant", + text="Kubernetes deployment", + start=0, + end=10, + ) + + +def test_search_uses_unicode_casefold() -> None: + index = ConversationSearchIndex( + lambda _: [{"role": "assistant", "content": "Straße"}] + ) + + assert index.filter([_conversation("one", "Travel")], "STRASSE") + assert index.get_match("one") == SearchMatch( + source="message", + role="assistant", + text="Straße", + start=0, + end=6, + ) + + +def test_search_fragment_precomputes_normalized_text_and_sparse_offsets() -> None: + ascii_fragment = SearchFragment(role="user", text="ALPHA beta") + unicode_fragment = SearchFragment(role="assistant", text="ß Straße") + + assert ascii_fragment.normalized_text == "alpha beta" + assert ascii_fragment._extra_normalized_offsets is None + assert unicode_fragment.normalized_text == "ss strasse" + assert unicode_fragment._extra_normalized_offsets is not None + assert len(unicode_fragment._extra_normalized_offsets) == 2 + assert unicode_fragment.find_span("sse") == (6, 8) + + +def test_search_fragment_does_not_renormalize_text_for_each_query() -> None: + class CountingText(str): + casefold_calls = 0 + + def casefold(self) -> str: + self.casefold_calls += 1 + return super().casefold() + + text = CountingText("Alpha beta") + fragment = SearchFragment(role="user", text=text) + + assert fragment.find_span("alpha") == (0, 5) + assert fragment.find_span("beta") == (6, 10) + assert fragment.find_span("missing") is None + assert text.casefold_calls == 1 + + +def test_match_snippet_collapses_whitespace_and_adds_ellipses() -> None: + text = f"{'a' * 80}\nmatched\t{'b' * 80}" + match = SearchMatch( + source="message", + role="user", + text=text, + start=81, + end=88, + ) + + snippet = match.snippet(context_chars=20) + + assert snippet.before == "a" * 19 + " " + assert snippet.matched == "matched" + assert snippet.after == " " + "b" * 19 + assert snippet.has_leading_ellipsis is True + assert snippet.has_trailing_ellipsis is True + + +def test_ignores_internal_non_text_and_non_chat_content() -> None: + messages = [ + {"role": "system", "content": "system-secret"}, + {"role": "tool", "content": "tool-secret"}, + {"role": "user", "content": "Memories related to the user request: memory"}, + {"role": "user", "content": "Content of notes.txt: injected-file"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "input": {"query": "tool-query"}}], + }, + ] + index = ConversationSearchIndex(lambda _: messages) + conversations = [_conversation("one", "Normal title")] + + for query in [ + "system-secret", + "tool-secret", + "memory", + "injected-file", + "tool-query", + ]: + assert index.filter(conversations, query) == [] + + +def test_history_is_loaded_once_per_browser_session() -> None: + calls = 0 + + def load_history(_: str) -> list[dict[str, Any]]: + nonlocal calls + calls += 1 + return [{"role": "user", "content": "alpha beta"}] + + index = ConversationSearchIndex(load_history) + conversations = [_conversation("one", "Unrelated")] + + assert index.filter(conversations, "alpha") + assert index.filter(conversations, "beta") + assert index.filter(conversations, "missing") == [] + assert calls == 1 + + +def test_loader_failure_is_cached_and_title_search_still_works() -> None: + calls = 0 + + def load_history(_: str) -> list[dict[str, Any]]: + nonlocal calls + calls += 1 + raise OSError("unreadable") + + index = ConversationSearchIndex(load_history) + conversations = [_conversation("one", "Readable title")] + + assert index.filter(conversations, "message") == [] + assert index.filter(conversations, "another") == [] + assert index.filter(conversations, "READABLE") == conversations + assert calls == 1 + + +def test_message_fragment_extractor_handles_malformed_content() -> None: + def extract_text(message: Any) -> list[str]: + return [ + fragment.text for fragment in iter_searchable_message_fragments(message) + ] + + assert extract_text({"role": "user", "content": "visible"}) == ["visible"] + assert extract_text(None) == [] + assert extract_text({"role": "user", "content": 42}) == [] + assert ( + extract_text( + { + "role": "assistant", + "content": [None, {"type": "text", "text": 42}], + } + ) + == [] + ) + + +def test_browser_ui_full_text_filter_and_filtered_delete() -> None: + histories = { + "one": [{"role": "user", "content": "alpha"}], + "two": [{"role": "assistant", "content": "beta"}], + } + ui = ConversationBrowserUI( + Console(width=120, height=40), + get_conversation_history=histories.get, + ) + ui.set_conversations( + [_conversation("one", "First"), _conversation("two", "Second")] + ) + + ui.update_search_query("beta") + assert [c["id"] for c in ui.conversations] == ["two"] + + ui.remove_conversations([0]) + ui.exit_search_mode(clear_filter=True) + assert [c["id"] for c in ui.conversations] == ["one"] + + +def test_browser_ui_renders_role_context_and_highlight_for_message_match() -> None: + ui = ConversationBrowserUI( + Console(width=120, height=40), + get_conversation_history=lambda _: [ + { + "role": "assistant", + "content": "Use PostgreSQL for durable storage", + } + ], + ) + ui.set_conversations([_conversation("one", "Database advice")]) + + ui.update_search_query("postgresql") + match = ui._search_index.get_match("one") + assert match is not None + + preview_lines = ui._create_search_match_preview(match) + rendered_match = preview_lines[-1] + + assert rendered_match.plain == "Assistant: Use PostgreSQL for durable storage" + highlighted_spans = [ + span + for span in rendered_match.spans + if str(span.style) == "bold black on yellow" + ] + assert len(highlighted_spans) == 1 + highlight = highlighted_spans[0] + assert rendered_match.plain[highlight.start : highlight.end] == "PostgreSQL" + + +def test_browser_ui_renders_title_match_without_loading_history() -> None: + calls = 0 + + def load_history(_: str) -> list[dict[str, Any]]: + nonlocal calls + calls += 1 + return [] + + ui = ConversationBrowserUI( + Console(width=120, height=40), + get_conversation_history=load_history, + ) + ui.set_conversations([_conversation("one", "Database advice")]) + + ui.update_search_query("database") + match = ui._search_index.get_match("one") + + assert match is not None + assert ui._create_search_match_preview(match)[-1].plain == "Title: Database advice" + assert calls == 0