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

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

2 

3from __future__ import annotations 

4 

5from litestar import delete, get, post, put 

6from litestar.exceptions import HTTPException 

7from litestar.params import Parameter 

8from litestar.response import Stream 

9from pydantic import BaseModel 

10 

11from lilbee.modelhub.role_validator import TaskMismatchError 

12from lilbee.server import handlers 

13from lilbee.server.auth import read_only 

14from lilbee.server.handlers import ModelsResponse, format_task_mismatch 

15from lilbee.server.models import ( 

16 ExternalModelsResponse, 

17 ModelsCatalogResponse, 

18 ModelsDeleteResponse, 

19 ModelsInstalledResponse, 

20 ModelsShowResponse, 

21 SetModelRequest, 

22 SetModelResponse, 

23) 

24 

25 

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

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

28 if isinstance(exc, TaskMismatchError): 

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

30 return str(exc) 

31 

32 

33class PullRequest(BaseModel): 

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

35 

36 model: str 

37 source: str = "native" 

38 allow_unsupported: bool = False 

39 

40 

41@get("/api/models") 

42@read_only 

43async def models_list_route() -> ModelsResponse: 

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

45 return await handlers.list_models() 

46 

47 

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

49@read_only 

50async def models_external_route() -> ExternalModelsResponse: 

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

52 return await handlers.list_external_models() 

53 

54 

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

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

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

58 try: 

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

60 except ValueError as exc: 

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

62 

63 

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

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

66 """Switch the active embedding model.""" 

67 try: 

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

69 except ValueError as exc: 

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

71 

72 

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

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

75 """Switch the active vision model for scanned PDF OCR. Empty disables OCR.""" 

76 try: 

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

78 except ValueError as exc: 

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

80 

81 

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

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

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

85 try: 

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

87 except ValueError as exc: 

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

89 

90 

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

92@read_only 

93async def models_catalog_route( 

94 task: str | None = Parameter(query="task", default=None), 

95 search: str = Parameter(query="search", default=""), 

96 size: str | None = Parameter(query="size", default=None), 

97 featured: bool | None = Parameter(query="featured", default=None), 

98 sort: str = Parameter(query="sort", default="featured"), 

99 limit: int = Parameter(query="limit", default=20, le=1000), 

100 offset: int = Parameter(query="offset", default=0, ge=0), 

101) -> ModelsCatalogResponse: 

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

103 try: 

104 return await handlers.models_catalog( 

105 task=task, 

106 search=search, 

107 size=size, 

108 featured=featured, 

109 sort=sort, 

110 limit=limit, 

111 offset=offset, 

112 ) 

113 except ValueError as exc: 

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

115 

116 

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

118@read_only 

119async def models_installed_route() -> ModelsInstalledResponse: 

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

121 return await handlers.models_installed() 

122 

123 

124@post("/api/models/pull") 

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

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

127 return Stream( 

128 handlers.models_pull( 

129 data.model, source=data.source, allow_unsupported=data.allow_unsupported 

130 ), 

131 media_type="text/event-stream", 

132 ) 

133 

134 

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

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

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

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

139 

140 

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

142async def models_delete_route(model: str, source: str = "native") -> ModelsDeleteResponse: 

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

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