Coverage for src / lilbee / server / app.py: 100%

1"""Litestar application factory: imports routes from modules, creates app with lifespan.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from collections.abc import AsyncIterator 

7from contextlib import asynccontextmanager 

8 

9from litestar import Litestar 

10from litestar.config.cors import CORSConfig 

11from litestar.middleware.base import DefineMiddleware 

12from litestar.openapi import OpenAPIConfig 

13 

14from lilbee.app.services import get_services 

15from lilbee.app.version import get_version 

16from lilbee.core.config import cfg 

17from lilbee.providers.sdk_llm_provider import inject_provider_keys 

18from lilbee.server.auth import AuthMiddleware, session_manager 

19from lilbee.server.routes.crawl import crawl_route 

20from lilbee.server.routes.documents import ( 

21 add_route, 

22 documents_list_route, 

23 documents_remove_route, 

24 export_route, 

25 import_route, 

26 sync_route, 

27) 

28from lilbee.server.routes.general import ( 

29 config_defaults_route, 

30 config_route, 

31 config_update_route, 

32 health_route, 

33 source_content_route, 

34 status_route, 

35) 

36from lilbee.server.routes.memory import ( 

37 memories_list_route, 

38 memories_remember_route, 

39 memories_remove_route, 

40 memories_update_route, 

41) 

42from lilbee.server.routes.models import ( 

43 models_catalog_route, 

44 models_delete_route, 

45 models_external_route, 

46 models_installed_route, 

47 models_list_route, 

48 models_pull_route, 

49 models_set_chat_route, 

50 models_set_embedding_route, 

51 models_set_reranker_route, 

52 models_set_vision_route, 

53 models_show_route, 

54) 

55from lilbee.server.routes.search import ( 

56 ask_route, 

57 ask_stream_route, 

58 chat_route, 

59 chat_stream_route, 

60 search_route, 

61) 

62from lilbee.server.routes.setup import ( 

63 setup_crawler_route, 

64 setup_crawler_status_route, 

65) 

66from lilbee.server.wiki import ( 

67 wiki_build_route, 

68 wiki_citations_reverse_route, 

69 wiki_draft_accept_route, 

70 wiki_draft_diff_route, 

71 wiki_draft_reject_route, 

72 wiki_drafts_route, 

73 wiki_lint_route, 

74 wiki_list_route, 

75 wiki_prune_route, 

76 wiki_read_route, 

77 wiki_status_route, 

78 wiki_synthesize_route, 

79 wiki_update_route, 

80) 

81 

82log = logging.getLogger(__name__) 

83 

84 

85@asynccontextmanager 

86async def _lifespan(app: Litestar) -> AsyncIterator[None]: 

87 """Pre-load LLM provider and embedding model on server startup.""" 

88 session_manager.load_or_generate() 

89 

90 inject_provider_keys() 

91 

92 try: 

93 get_services() # pre-load all services (provider, embedder, etc.) 

94 log.info("LLM provider pre-loaded") 

95 except Exception: 

96 log.warning("Failed to pre-load LLM provider", exc_info=True) 

97 try: 

98 get_services().embedder.validate_model() 

99 log.info("Embedding model validated") 

100 except Exception: 

101 log.warning("Failed to validate embedding model", exc_info=True) 

102 try: 

103 yield 

104 finally: 

105 session_manager.cleanup() 

106 

107 

108def create_app() -> Litestar: 

109 """Create the Litestar application instance.""" 

110 cors = CORSConfig( 

111 allow_origins=cfg.cors_origins, 

112 allow_origin_regex=cfg.cors_origin_regex, 

113 allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH"], 

114 allow_headers=["Content-Type", "Authorization"], 

115 ) 

116 return Litestar( 

117 lifespan=[_lifespan], 

118 middleware=[DefineMiddleware(AuthMiddleware)], 

119 route_handlers=[ 

120 health_route, 

121 status_route, 

122 config_route, 

123 config_defaults_route, 

124 config_update_route, 

125 source_content_route, 

126 search_route, 

127 ask_route, 

128 ask_stream_route, 

129 chat_route, 

130 chat_stream_route, 

131 sync_route, 

132 add_route, 

133 models_list_route, 

134 models_external_route, 

135 models_set_chat_route, 

136 models_set_embedding_route, 

137 models_set_vision_route, 

138 models_set_reranker_route, 

139 models_catalog_route, 

140 models_installed_route, 

141 models_pull_route, 

142 models_show_route, 

143 models_delete_route, 

144 documents_list_route, 

145 documents_remove_route, 

146 memories_list_route, 

147 memories_remember_route, 

148 memories_update_route, 

149 memories_remove_route, 

150 export_route, 

151 import_route, 

152 crawl_route, 

153 setup_crawler_route, 

154 setup_crawler_status_route, 

155 wiki_list_route, 

156 wiki_read_route, 

157 wiki_drafts_route, 

158 wiki_draft_diff_route, 

159 wiki_draft_accept_route, 

160 wiki_draft_reject_route, 

161 wiki_citations_reverse_route, 

162 wiki_lint_route, 

163 wiki_prune_route, 

164 wiki_build_route, 

165 wiki_update_route, 

166 wiki_status_route, 

167 wiki_synthesize_route, 

168 ], 

169 request_max_body_size=10 * 1024 * 1024, 

170 cors_config=cors, 

171 openapi_config=OpenAPIConfig( 

172 title="lilbee", 

173 description="Local knowledge base REST API", 

174 version=get_version(), 

175 path="/schema", 

176 ), 

177 )