Coverage for src/lilbee/providers/roles.py: 100%

1"""Engine-neutral role identifiers shared across the provider stack. 

2 

3``WorkerRole`` names the four inference roles a local engine serves; the fleet 

4maps each to one llama-server instance. ``OcrBackend`` names the PDF-OCR paths. 

5``EngineBackend`` names the compute backend the engine serves those roles on. 

6These outlive any particular engine, so they live here rather than inside an 

7engine-specific module. 

8""" 

9 

10from __future__ import annotations 

11 

12from dataclasses import dataclass 

13from enum import StrEnum 

14from typing import Literal 

15 

16 

17class WorkerRole(StrEnum): 

18 """Inference role identifier; addresses one llama-server in the fleet.""" 

19 

20 EMBED = "embed" 

21 RERANK = "rerank" 

22 CHAT = "chat" 

23 VISION = "vision" 

24 

25 

26class EngineBackend(StrEnum): 

27 """The compute backend the engine selected on this host. 

28 

29 ``CPU`` and ``UNKNOWN`` are separate answers on purpose. An empty device list 

30 has two causes a client cannot tell apart: a host that has no usable GPU, and 

31 a host whose device probe never answered. Only the first is CPU. Reporting 

32 ``CPU`` for the second misidentifies a CUDA or Metal host in the diagnostics 

33 a user pastes into a bug report. 

34 """ 

35 

36 CUDA = "cuda" 

37 ROCM = "rocm" 

38 METAL = "metal" 

39 SYCL = "sycl" 

40 VULKAN = "vulkan" 

41 CPU = "cpu" 

42 UNKNOWN = "unknown" 

43 

44 

45class Phase(StrEnum): 

46 """A run phase whose roles are loaded together on demand. 

47 

48 Ingest OCRs and embeds (vision + embed); a query embeds, reranks, and generates 

49 (embed + rerank + chat). Roles sharing no phase are never co-resident, so on a 

50 tight host they may share one swap group instead of both reserving VRAM. 

51 """ 

52 

53 INGEST = "ingest" 

54 QUERY = "query" 

55 

56 

57class RerankMode(StrEnum): 

58 """Resolved reranker serving mode for one RERANK server. 

59 

60 ``CROSS_ENCODER`` serves an encoder GGUF with rank-pooling embeddings; 

61 ``LLM`` serves a decoder GGUF generatively and scores yes/no logprobs. 

62 """ 

63 

64 CROSS_ENCODER = "cross_encoder" 

65 LLM = "llm" 

66 

67 

68@dataclass(frozen=True) 

69class RoleInfo: 

70 """The per-role knowledge the fleet needs to configure one llama-server. 

71 

72 One row per ``WorkerRole``, so adding a role is a single registry entry and the 

73 scattered planning/placement/replica tuples all derive from here. 

74 """ 

75 

76 role: WorkerRole 

77 config_field: str # the cfg ``*_model`` field whose value this role serves 

78 replicated: bool # runs N data-parallel replicas (embed/vision) 

79 replica_knob: str | None # cfg int field scaling replicas, None when not replicated 

80 offload_all_layers: bool # loader offloads every layer, ignoring cfg.n_gpu_layers 

81 flash_attn: bool # runs with flash attention (chat/vision) 

82 pooled: bool # pooled single-slot search role (embed/cross-encoder rerank) 

83 placement_rank: int # placement order; the elastic chat model is charged last 

84 phases: frozenset[Phase] # run phases that load this role (co-residency model) 

85 

86 

87ROLE_REGISTRY: dict[WorkerRole, RoleInfo] = { 

88 WorkerRole.CHAT: RoleInfo( 

89 role=WorkerRole.CHAT, 

90 config_field="chat_model", 

91 replicated=False, 

92 replica_knob=None, 

93 offload_all_layers=False, 

94 flash_attn=True, 

95 pooled=False, 

96 placement_rank=2, 

97 phases=frozenset({Phase.QUERY}), 

98 ), 

99 WorkerRole.EMBED: RoleInfo( 

100 role=WorkerRole.EMBED, 

101 config_field="embedding_model", 

102 replicated=True, 

103 replica_knob="embed_replicas", 

104 offload_all_layers=True, 

105 flash_attn=False, 

106 pooled=True, 

107 placement_rank=0, 

108 phases=frozenset({Phase.INGEST, Phase.QUERY}), 

109 ), 

110 WorkerRole.RERANK: RoleInfo( 

111 role=WorkerRole.RERANK, 

112 config_field="reranker_model", 

113 replicated=False, 

114 replica_knob=None, 

115 offload_all_layers=True, 

116 flash_attn=False, 

117 pooled=True, 

118 placement_rank=0, 

119 phases=frozenset({Phase.QUERY}), 

120 ), 

121 WorkerRole.VISION: RoleInfo( 

122 role=WorkerRole.VISION, 

123 config_field="vision_model", 

124 replicated=True, 

125 replica_knob="vision_replicas", 

126 offload_all_layers=True, 

127 flash_attn=True, 

128 pooled=False, 

129 placement_rank=1, 

130 phases=frozenset({Phase.INGEST}), 

131 ), 

132} 

133"""Single source of truth for per-role fleet configuration, ordered chat/embed/rerank/vision.""" 

134 

135 

136MODEL_FIELD_TO_ROLE: dict[str, WorkerRole] = { 

137 info.config_field: role for role, info in ROLE_REGISTRY.items() 

138} 

139"""Config model-role field name -> the worker whose server serves it. 

140 

141A model-role setting change reloads just that role's server (off-thread) rather 

142than dropping the whole fleet, so unrelated roles keep serving uninterrupted. 

143""" 

144 

145 

146MODEL_ROLE_FIELDS: frozenset[str] = frozenset(MODEL_FIELD_TO_ROLE) 

147"""The cfg ``*_model`` field names, as a set (settings overlay + reload routing).""" 

148 

149 

150ROLE_GATE_FIELD_TO_ROLE: dict[str, WorkerRole] = {"enable_ocr": WorkerRole.VISION} 

151"""Non-model settings that decide whether a role is planned; a change reloads that role.""" 

152 

153 

154REPLICATED_ROLES: tuple[WorkerRole, ...] = tuple( 

155 role for role, info in ROLE_REGISTRY.items() if info.replicated 

156) 

157"""Roles whose ``*_replicas`` knob scales data-parallel instances; others run one.""" 

158 

159 

160OcrBackend = Literal["vision"] 

161"""PDF-OCR backends routed to the engine. Tesseract runs inline, not on a server.""" 

162 

163 

164def configured_model_message(role: WorkerRole, configured: str, requested: str) -> str: 

165 """User-facing rejection for a per-call model that differs from the configured one.""" 

166 return ( 

167 f"This engine serves the configured {role} model ({configured}). " 

168 f"To use {requested!r}, set it as the {role} model in lilbee settings " 

169 f"(TUI /settings), then retry; the engine reloads automatically." 

170 )