Coverage for src / lilbee / cli / tui / screens / catalog_grouping.py: 100%

1"""Row-grouping helpers and the GridSection container for CatalogScreen.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6 

7from lilbee.catalog.types import ModelTask 

8from lilbee.cli.tui import messages as msg 

9from lilbee.cli.tui.screens.catalog_utils import ( 

10 CatalogRow, 

11 CatalogRowKind, 

12 FrontierCatalogRow, 

13 LocalCatalogRow, 

14) 

15from lilbee.cli.tui.widgets.model_list import ModelListSection 

16 

17 

18@dataclass 

19class GridSection: 

20 """A named group of rows for the grid view.""" 

21 

22 heading: str 

23 rows: list[CatalogRow] 

24 

25 

26TASK_BUCKET_ORDER = (ModelTask.CHAT, ModelTask.EMBEDDING, ModelTask.VISION, ModelTask.RERANK) 

27PICKS_SECTION_HEADING = "★ Picks" 

28 

29 

30def row_cache_signature(row: CatalogRow) -> tuple[str, bool]: 

31 """Pair (name, installed-flag) for the per-tab cache key. 

32 

33 Frontier rows don't carry an ``installed`` field; they're keyed as 

34 if installed=False since each frontier entry is provider-managed 

35 rather than on-disk. 

36 """ 

37 if row.kind == CatalogRowKind.FRONTIER: 

38 return (row.name, False) 

39 return (row.name, row.installed) 

40 

41 

42def for_you_sort_key(row: LocalCatalogRow) -> tuple[int, str]: 

43 """Rank Discover 'For You' rows: best fit first, then alphabetical. 

44 

45 Fit rank: FITS=0, TIGHT=1, WONT_RUN=2, no chip=3. Featured-only 

46 callers already filtered, so featured isn't in the key. 

47 """ 

48 from lilbee.runtime.hardware import FitLevel 

49 

50 if row.fit is None: 

51 rank = 3 

52 elif row.fit.level is FitLevel.FITS: 

53 rank = 0 

54 elif row.fit.level is FitLevel.TIGHT: 

55 rank = 1 

56 else: 

57 rank = 2 

58 return (rank, row.name.lower()) 

59 

60 

61def group_frontier_rows( 

62 frontier_rows: list[FrontierCatalogRow], 

63) -> list[ModelListSection]: 

64 """Group frontier rows into provider-headed sections. 

65 

66 Section order follows :data:`PROVIDER_KEYS` (the canonical display 

67 order); providers absent from PROVIDER_KEYS land at the tail in 

68 alphabetical order. Rows within each section are alphabetical. 

69 """ 

70 if not frontier_rows: 

71 return [] 

72 from lilbee.providers.sdk_backend import PROVIDER_KEYS 

73 

74 per_provider: dict[str, list[FrontierCatalogRow]] = {} 

75 for row in frontier_rows: 

76 per_provider.setdefault(row.provider, []).append(row) 

77 canonical_order = [label for _, _, _, label in PROVIDER_KEYS] 

78 ordered = [p for p in canonical_order if p in per_provider] 

79 extras = sorted(set(per_provider) - set(canonical_order)) 

80 sections: list[ModelListSection] = [] 

81 for provider in [*ordered, *extras]: 

82 rows = sorted(per_provider[provider], key=lambda r: r.name.lower()) 

83 sections.append(ModelListSection(heading=provider, rows=list(rows))) 

84 return sections 

85 

86 

87def group_task_rows_with_picks( 

88 task_rows: list[LocalCatalogRow], task_label: str 

89) -> list[GridSection]: 

90 """Per-tab grouping: ★ Picks pinned, then Installed, then the rest. 

91 

92 Lifts featured rows out of their task bucket into a dedicated pinned 

93 section at the top of the tab. Today's behavior interleaved them at 

94 the top of the task bucket; the redesign treats curation as its own 

95 layer so the eye lands on Picks first instead of having to scan past 

96 them to find non-featured rows. 

97 

98 Pre-condition: caller has already filtered ``task_rows`` to a single 

99 task (the active per-task tab). 

100 """ 

101 picks: list[CatalogRow] = [] 

102 installed: list[CatalogRow] = [] 

103 others: list[CatalogRow] = [] 

104 for row in task_rows: 

105 if row.featured: 

106 picks.append(row) 

107 elif row.installed: 

108 installed.append(row) 

109 else: 

110 others.append(row) 

111 return [ 

112 GridSection(PICKS_SECTION_HEADING, picks), 

113 GridSection(msg.HEADING_INSTALLED, installed), 

114 GridSection(task_label, others), 

115 ] 

116 

117 

118def group_rows_for_grid(local_rows: list[LocalCatalogRow]) -> list[GridSection]: 

119 """Group local rows into sections for the grid view. 

120 

121 Layout: Installed first, then one section per task. Featured rows live 

122 at the top of their task section (recognizable by the ``pick`` pill); 

123 no separate "Our picks" bucket so the catalog reads as a single 

124 task-organized list. 

125 """ 

126 installed: list[CatalogRow] = [] 

127 by_task: dict[str, list[CatalogRow]] = {task: [] for task in TASK_BUCKET_ORDER} 

128 extras: dict[str, list[CatalogRow]] = {} 

129 for row in local_rows: 

130 if row.installed: 

131 installed.append(row) 

132 continue 

133 bucket = by_task.get(row.task) 

134 if bucket is not None: 

135 bucket.append(row) 

136 else: 

137 extras.setdefault(row.task, []).append(row) 

138 # Within each task bucket: featured first (preserving their input order), 

139 # then the rest in their incoming order. Stable so HF rank from the API 

140 # is preserved among non-featured rows. 

141 for bucket in by_task.values(): 

142 bucket.sort(key=lambda r: not getattr(r, "featured", False)) 

143 for bucket in extras.values(): 

144 bucket.sort(key=lambda r: not getattr(r, "featured", False)) 

145 return [ 

146 GridSection(msg.HEADING_INSTALLED, installed), 

147 *[GridSection(task.capitalize(), by_task[task]) for task in TASK_BUCKET_ORDER], 

148 *[GridSection(task.capitalize(), extras[task]) for task in extras], 

149 ]