Coverage for src/local_deep_research/web_search_engines/engine_groups.py: 100%

1"""Grouping and ordering of search engines in the selector UI. 

2 

3Single source of truth for which band a search engine falls into and the order 

4the bands appear in. The available-search-engines API attaches each engine's 

5band to the response and the frontend renders one header per band, so the 

6bucketing logic lives here only and is never duplicated in JavaScript. 

7""" 

8 

9from dataclasses import dataclass 

10 

11 

12@dataclass(frozen=True) 

13class SearchEngineGroup: 

14 """A band in the search-engine selector dropdown.""" 

15 

16 key: str # stable identifier, e.g. "academic" 

17 label: str # header text shown in the dropdown 

18 icon: str # emoji associated with the band 

19 

20 

21# Bands in display order, top to bottom; the list index IS the sort order 

22# (lower sorts higher). "favorites" is an overlay band: a starred engine is 

23# shown here regardless of its category — see effective_group(). 

24SEARCH_ENGINE_GROUPS: tuple[SearchEngineGroup, ...] = ( 

25 SearchEngineGroup("favorites", "Favorites", "⭐"), 

26 SearchEngineGroup("collections", "Collections", "📁"), 

27 SearchEngineGroup("academic", "Academic", "🔬"), 

28 SearchEngineGroup("local_rag", "Local RAG", "📂"), 

29 SearchEngineGroup("books", "Books", "📚"), 

30 SearchEngineGroup("code", "Code", "💻"), 

31 SearchEngineGroup("news", "News", "📰"), 

32 SearchEngineGroup("no_api_key", "No API key", "🌐"), 

33 SearchEngineGroup("api_key", "API key", "🔑"), 

34) 

35 

36# Not a secret — the band's identifier string; matches SEARCH_ENGINE_GROUPS[0]. 

37FAVORITES_GROUP_KEY = "favorites" # gitleaks:allow 

38 

39_GROUP_BY_KEY: dict[str, SearchEngineGroup] = { 

40 group.key: group for group in SEARCH_ENGINE_GROUPS 

41} 

42_GROUP_ORDER: dict[str, int] = { 

43 group.key: index for index, group in enumerate(SEARCH_ENGINE_GROUPS) 

44} 

45 

46# Maps the category label produced by 

47# settings_routes._get_engine_icon_and_category() to a band. Any category not 

48# listed here (e.g. "Web Search", "Search") falls through to the API-key split. 

49# Keep these strings in sync with that function. 

50_CATEGORY_TO_GROUP: dict[str, str] = { 

51 "Scientific": "academic", 

52 "Local RAG": "local_rag", 

53 "Books": "books", 

54 "Code": "code", 

55 "News": "news", 

56} 

57 

58 

59def is_collection_engine(engine_id: str) -> bool: 

60 """True for the document-collection pseudo-engines: the aggregate 

61 "Search All Collections" entry (``library``) and the per-collection 

62 ``collection_<id>`` entries.""" 

63 return engine_id == "library" or engine_id.startswith("collection_") 

64 

65 

66def classify_engine_group( 

67 engine_id: str, category: str, requires_api_key: bool 

68) -> str: 

69 """Return the base band key for an engine, ignoring favorite status. 

70 

71 Document collections come first, then engines with a distinct category 

72 (academic / local / books / code / news); everything else (generic web 

73 search) is split only by whether it needs an API key. 

74 """ 

75 if is_collection_engine(engine_id): 

76 return "collections" 

77 category_group = _CATEGORY_TO_GROUP.get(category) 

78 if category_group is not None: 

79 return category_group 

80 return "api_key" if requires_api_key else "no_api_key" 

81 

82 

83def effective_group(base_group_key: str, is_favorite: bool) -> str: 

84 """The band an engine is actually shown in: the Favorites overlay wins over 

85 the engine's base category so a starred engine floats to the top band.""" 

86 return FAVORITES_GROUP_KEY if is_favorite else base_group_key 

87 

88 

89def group_order(group_key: str) -> int: 

90 """Sort index for a band (lower sorts higher); unknown keys sort last.""" 

91 return _GROUP_ORDER.get(group_key, len(SEARCH_ENGINE_GROUPS)) 

92 

93 

94def group_label(group_key: str) -> str: 

95 """Header text for a band; falls back to the key if unknown.""" 

96 group = _GROUP_BY_KEY.get(group_key) 

97 return group.label if group is not None else group_key 

98 

99 

100def group_icon(group_key: str) -> str: 

101 """Emoji for a band; empty string if unknown.""" 

102 group = _GROUP_BY_KEY.get(group_key) 

103 return group.icon if group is not None else ""