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 

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() -> ModelsInstalledResponse: 

134 """List installed models with their source (native or remote).""" 

135 return await handlers.models_installed() 

136 

137 

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

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

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

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

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

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

144 try: 

145 await handlers.enforce_pull_arch_compat( 

146 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

147 ) 

148 except ValueError as exc: 

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

150 return Stream( 

151 handlers.models_pull( 

152 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

153 ), 

154 media_type=SSE_MEDIA_TYPE, 

155 ) 

156 

157 

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

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

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

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

162 

163 

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

165async def models_delete_route( 

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

167) -> ModelsDeleteResponse: 

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

169 try: 

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

171 except ValueError as exc: 

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