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 disables OCR.""" 

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 sort: FromQuery[str] = "featured", 

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

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

106) -> ModelsCatalogResponse: 

107 """Browse the model catalog with optional filters.""" 

108 try: 

109 return await handlers.models_catalog( 

110 task=task, 

111 search=search, 

112 size=size, 

113 installed=installed, 

114 featured=featured, 

115 sort=sort, 

116 limit=limit, 

117 offset=offset, 

118 ) 

119 except ValueError as exc: 

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

121 

122 

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

124async def models_installed_route() -> ModelsInstalledResponse: 

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

126 return await handlers.models_installed() 

127 

128 

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

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

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

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

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

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

135 try: 

136 await handlers.enforce_pull_arch_compat( 

137 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

138 ) 

139 except ValueError as exc: 

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

141 return Stream( 

142 handlers.models_pull( 

143 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

144 ), 

145 media_type=SSE_MEDIA_TYPE, 

146 ) 

147 

148 

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

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

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

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

153 

154 

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

156async def models_delete_route( 

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

158) -> ModelsDeleteResponse: 

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

160 try: 

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

162 except ValueError as exc: 

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