Coverage for src/lilbee/server/routes/models.py: 100%

1"""Model management route handlers: catalog, installed, pull, show, delete, set. 

2 

3Every route needs the token: the reads describe what the user has installed 

4and what their machine fits, which is host inventory. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import Annotated 

10 

11from litestar import delete, get, post, put 

12from litestar.exceptions import HTTPException 

13from litestar.params import FromPath, FromQuery, QueryParameter 

14from litestar.response import Stream 

15from pydantic import BaseModel 

16 

17from lilbee.catalog.types import ModelSource, ModelTask 

18from lilbee.modelhub.role_validator import TaskMismatchError 

19from lilbee.server import handlers 

20from lilbee.server.handlers import ModelsResponse, format_task_mismatch 

21from lilbee.server.handlers.sse import SSE_MEDIA_TYPE 

22from lilbee.server.models import ( 

23 ExternalModelsResponse, 

24 ModelsCatalogResponse, 

25 ModelsDeleteResponse, 

26 ModelsInstalledResponse, 

27 ModelsShowResponse, 

28 SetModelRequest, 

29 SetModelResponse, 

30) 

31 

32 

33def _task_mismatch_detail(exc: ValueError) -> str: 

34 """Format a 422 detail string, expanding TaskMismatchError into HTTP guidance.""" 

35 if isinstance(exc, TaskMismatchError): 

36 return format_task_mismatch(exc.ref, exc.entry_task, exc.expected_task) 

37 return str(exc) 

38 

39 

40class PullRequest(BaseModel): 

41 """Request body for /api/models/pull.""" 

42 

43 model: str 

44 source: str = ModelSource.NATIVE.value 

45 allow_unsupported: bool = False 

46 

47 

48@get("/api/models") 

49async def models_list_route() -> ModelsResponse: 

50 """Available chat, embedding, vision, and reranker models.""" 

51 return await handlers.list_models() 

52 

53 

54@get("/api/models/external") 

55async def models_external_route() -> ExternalModelsResponse: 

56 """Discover models available from the configured external provider.""" 

57 return await handlers.list_external_models() 

58 

59 

60@put("/api/models/chat") 

61async def models_set_chat_route(data: SetModelRequest) -> SetModelResponse: 

62 """Switch the active chat model used for RAG answers.""" 

63 try: 

64 return await handlers.set_chat_model(model=data.model) 

65 except ValueError as exc: 

66 raise HTTPException(status_code=422, detail=_task_mismatch_detail(exc)) from exc 

67 

68 

69@put("/api/models/embedding") 

70async def models_set_embedding_route(data: SetModelRequest) -> SetModelResponse: 

71 """Switch the active embedding model.""" 

72 try: 

73 return await handlers.set_embedding_model(model=data.model) 

74 except ValueError as exc: 

75 raise HTTPException(status_code=422, detail=_task_mismatch_detail(exc)) from exc 

76 

77 

78@put("/api/models/vision") 

79async def models_set_vision_route(data: SetModelRequest) -> SetModelResponse: 

80 """Switch the active vision model for scanned PDF OCR. Empty falls back to Tesseract.""" 

81 try: 

82 return await handlers.set_vision_model(model=data.model) 

83 except ValueError as exc: 

84 raise HTTPException(status_code=422, detail=_task_mismatch_detail(exc)) from exc 

85 

86 

87@put("/api/models/reranker") 

88async def models_set_reranker_route(data: SetModelRequest) -> SetModelResponse: 

89 """Switch the active reranker model. Empty disables reranking.""" 

90 try: 

91 return await handlers.set_reranker_model(model=data.model) 

92 except ValueError as exc: 

93 raise HTTPException(status_code=422, detail=_task_mismatch_detail(exc)) from exc 

94 

95 

96@get("/api/models/catalog") 

97async def models_catalog_route( 

98 task: FromQuery[str | None] = None, 

99 search: FromQuery[str] = "", 

100 size: FromQuery[str | None] = None, 

101 installed: FromQuery[bool | None] = None, 

102 featured: FromQuery[bool | None] = None, 

103 max_fit: FromQuery[str | None] = None, 

104 sort: FromQuery[str] = "featured", 

105 limit: Annotated[int, QueryParameter(ge=1, le=1000)] = 20, 

106 offset: Annotated[int, QueryParameter(ge=0)] = 0, 

107 source: FromQuery[str | None] = None, 

108) -> ModelsCatalogResponse: 

109 """Browse the model catalog with optional filters, applied before paging. 

110 

111 ``max_fit`` is the worst hardware fit to return: ``fits``, ``tight``, or 

112 ``wont_run``. ``source`` keeps one row source, e.g. ``native`` for local 

113 GGUF rows only. 

114 """ 

115 try: 

116 return await handlers.models_catalog( 

117 task=task, 

118 search=search, 

119 size=size, 

120 installed=installed, 

121 featured=featured, 

122 max_fit=max_fit, 

123 sort=sort, 

124 limit=limit, 

125 offset=offset, 

126 source=source, 

127 ) 

128 except ValueError as exc: 

129 raise HTTPException(status_code=422, detail=str(exc)) from exc 

130 

131 

132@get("/api/models/installed") 

133async def models_installed_route(task: FromQuery[str | None] = None) -> ModelsInstalledResponse: 

134 """List installed models with source, task, size and display name; ``task`` keeps one role.""" 

135 try: 

136 parsed_task = ModelTask(task) if task else None 

137 except ValueError as exc: 

138 raise HTTPException(status_code=422, detail=str(exc)) from exc 

139 return await handlers.models_installed(task=parsed_task) 

140 

141 

142@post("/api/models/pull", media_type=SSE_MEDIA_TYPE) 

143async def models_pull_route(data: PullRequest) -> Stream: 

144 """Pull a model with streaming SSE progress events.""" 

145 # Validate before opening the stream so an unsupported arch is a real 409, not 

146 # an in-stream abort after the 200 SSE headers have flushed. An unknown source 

147 # is a client error (422), mirroring the catalog route, not a 500. 

148 try: 

149 await handlers.enforce_pull_arch_compat( 

150 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

151 ) 

152 except ValueError as exc: 

153 raise HTTPException(status_code=422, detail=str(exc)) from exc 

154 return Stream( 

155 handlers.models_pull( 

156 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

157 ), 

158 media_type=SSE_MEDIA_TYPE, 

159 ) 

160 

161 

162@post("/api/models/show") 

163async def models_show_route(data: SetModelRequest) -> ModelsShowResponse: 

164 """Get model metadata and parameter defaults.""" 

165 return await handlers.models_show(model=data.model) 

166 

167 

168@delete("/api/models/{model:str}", status_code=200) 

169async def models_delete_route( 

170 model: FromPath[str], source: FromQuery[str] = ModelSource.NATIVE.value 

171) -> ModelsDeleteResponse: 

172 """Delete a model from the specified source.""" 

173 try: 

174 return await handlers.models_delete(model, source=source) 

175 except ValueError as exc: 

176 raise HTTPException(status_code=422, detail=str(exc)) from exc