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.server import handlers 

12from lilbee.server.auth import read_only 

13from lilbee.server.handlers import ModelsResponse 

14from lilbee.server.models import ( 

15 ExternalModelsResponse, 

16 ModelsCatalogResponse, 

17 ModelsDeleteResponse, 

18 ModelsInstalledResponse, 

19 ModelsShowResponse, 

20 SetModelRequest, 

21 SetModelResponse, 

22) 

23 

24 

25class PullRequest(BaseModel): 

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

27 

28 model: str 

29 source: str = "native" 

30 

31 

32@get("/api/models") 

33@read_only 

34async def models_list_route() -> ModelsResponse: 

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

36 return await handlers.list_models() 

37 

38 

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

40@read_only 

41async def models_external_route() -> ExternalModelsResponse: 

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

43 return await handlers.list_external_models() 

44 

45 

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

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

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

49 try: 

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

51 except ValueError as exc: 

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

53 

54 

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

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

57 """Switch the active embedding model.""" 

58 try: 

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

60 except ValueError as exc: 

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

62 

63 

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

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

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

67 try: 

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

69 except ValueError as exc: 

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

71 

72 

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

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

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

76 try: 

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

78 except ValueError as exc: 

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

80 

81 

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

83@read_only 

84async def models_catalog_route( 

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

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

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

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

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

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

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

92) -> ModelsCatalogResponse: 

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

94 return await handlers.models_catalog( 

95 task=task, 

96 search=search, 

97 size=size, 

98 featured=featured, 

99 sort=sort, 

100 limit=limit, 

101 offset=offset, 

102 ) 

103 

104 

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

106@read_only 

107async def models_installed_route() -> ModelsInstalledResponse: 

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

109 return await handlers.models_installed() 

110 

111 

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

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

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

115 return Stream( 

116 handlers.models_pull(data.model, source=data.source), 

117 media_type="text/event-stream", 

118 ) 

119 

120 

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

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

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

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

125 

126 

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

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

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

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