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. 

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

6engine-specific module. 

7""" 

8 

9from __future__ import annotations 

10 

11from dataclasses import dataclass 

12from enum import StrEnum 

13from typing import Literal 

14 

15 

16class WorkerRole(StrEnum): 

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

18 

19 EMBED = "embed" 

20 RERANK = "rerank" 

21 CHAT = "chat" 

22 VISION = "vision" 

23 

24 

25class Phase(StrEnum): 

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

27 

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

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

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

31 """ 

32 

33 INGEST = "ingest" 

34 QUERY = "query" 

35 

36 

37class RerankMode(StrEnum): 

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

39 

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

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

42 """ 

43 

44 CROSS_ENCODER = "cross_encoder" 

45 LLM = "llm" 

46 

47 

48@dataclass(frozen=True) 

49class RoleInfo: 

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

51 

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

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

54 """ 

55 

56 role: WorkerRole 

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

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

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

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

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

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

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

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

65 

66 

67ROLE_REGISTRY: dict[WorkerRole, RoleInfo] = { 

68 WorkerRole.CHAT: RoleInfo( 

69 role=WorkerRole.CHAT, 

70 config_field="chat_model", 

71 replicated=False, 

72 replica_knob=None, 

73 offload_all_layers=False, 

74 flash_attn=True, 

75 pooled=False, 

76 placement_rank=2, 

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

78 ), 

79 WorkerRole.EMBED: RoleInfo( 

80 role=WorkerRole.EMBED, 

81 config_field="embedding_model", 

82 replicated=True, 

83 replica_knob="embed_replicas", 

84 offload_all_layers=True, 

85 flash_attn=False, 

86 pooled=True, 

87 placement_rank=0, 

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

89 ), 

90 WorkerRole.RERANK: RoleInfo( 

91 role=WorkerRole.RERANK, 

92 config_field="reranker_model", 

93 replicated=False, 

94 replica_knob=None, 

95 offload_all_layers=True, 

96 flash_attn=False, 

97 pooled=True, 

98 placement_rank=0, 

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

100 ), 

101 WorkerRole.VISION: RoleInfo( 

102 role=WorkerRole.VISION, 

103 config_field="vision_model", 

104 replicated=True, 

105 replica_knob="vision_replicas", 

106 offload_all_layers=True, 

107 flash_attn=True, 

108 pooled=False, 

109 placement_rank=1, 

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

111 ), 

112} 

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

114 

115 

116MODEL_FIELD_TO_ROLE: dict[str, WorkerRole] = { 

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

118} 

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

120 

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

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

123""" 

124 

125 

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

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

128 

129 

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

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

132) 

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

134 

135 

136OcrBackend = Literal["vision"] 

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

138 

139 

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

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

142 return ( 

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

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

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

146 )