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.cli.helpers import get_version 

15from lilbee.config import cfg 

16from lilbee.providers.sdk_llm_provider import inject_provider_keys 

17from lilbee.server.auth import AuthMiddleware, session_manager 

18from lilbee.server.routes.crawl import crawl_route 

19from lilbee.server.routes.documents import ( 

20 add_route, 

21 documents_list_route, 

22 documents_remove_route, 

23 sync_route, 

24) 

25from lilbee.server.routes.general import ( 

26 config_defaults_route, 

27 config_route, 

28 config_update_route, 

29 health_route, 

30 source_content_route, 

31 status_route, 

32) 

33from lilbee.server.routes.models import ( 

34 models_catalog_route, 

35 models_delete_route, 

36 models_external_route, 

37 models_installed_route, 

38 models_list_route, 

39 models_pull_route, 

40 models_set_chat_route, 

41 models_set_embedding_route, 

42 models_set_reranker_route, 

43 models_set_vision_route, 

44 models_show_route, 

45) 

46from lilbee.server.routes.search import ( 

47 ask_route, 

48 ask_stream_route, 

49 chat_route, 

50 chat_stream_route, 

51 search_route, 

52) 

53from lilbee.server.routes.setup import ( 

54 setup_crawler_route, 

55 setup_crawler_status_route, 

56) 

57from lilbee.server.wiki import ( 

58 wiki_build_route, 

59 wiki_citations_reverse_route, 

60 wiki_draft_accept_route, 

61 wiki_draft_diff_route, 

62 wiki_draft_reject_route, 

63 wiki_drafts_route, 

64 wiki_lint_route, 

65 wiki_list_route, 

66 wiki_prune_route, 

67 wiki_read_route, 

68 wiki_status_route, 

69 wiki_synthesize_route, 

70 wiki_update_route, 

71) 

72from lilbee.services import get_services 

73 

74log = logging.getLogger(__name__) 

75 

76 

77@asynccontextmanager 

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

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

80 session_manager.load_or_generate() 

81 

82 inject_provider_keys() 

83 

84 try: 

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

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

87 except Exception: 

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

89 try: 

90 get_services().embedder.validate_model() 

91 log.info("Embedding model validated") 

92 except Exception: 

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

94 try: 

95 yield 

96 finally: 

97 session_manager.cleanup() 

98 

99 

100def create_app() -> Litestar: 

101 """Create the Litestar application instance.""" 

102 cors = CORSConfig( 

103 allow_origins=cfg.cors_origins, 

104 allow_origin_regex=cfg.cors_origin_regex, 

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

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

107 ) 

108 return Litestar( 

109 lifespan=[_lifespan], 

110 middleware=[DefineMiddleware(AuthMiddleware)], 

111 route_handlers=[ 

112 health_route, 

113 status_route, 

114 config_route, 

115 config_defaults_route, 

116 config_update_route, 

117 source_content_route, 

118 search_route, 

119 ask_route, 

120 ask_stream_route, 

121 chat_route, 

122 chat_stream_route, 

123 sync_route, 

124 add_route, 

125 models_list_route, 

126 models_external_route, 

127 models_set_chat_route, 

128 models_set_embedding_route, 

129 models_set_vision_route, 

130 models_set_reranker_route, 

131 models_catalog_route, 

132 models_installed_route, 

133 models_pull_route, 

134 models_show_route, 

135 models_delete_route, 

136 documents_list_route, 

137 documents_remove_route, 

138 crawl_route, 

139 setup_crawler_route, 

140 setup_crawler_status_route, 

141 wiki_list_route, 

142 wiki_read_route, 

143 wiki_drafts_route, 

144 wiki_draft_diff_route, 

145 wiki_draft_accept_route, 

146 wiki_draft_reject_route, 

147 wiki_citations_reverse_route, 

148 wiki_lint_route, 

149 wiki_prune_route, 

150 wiki_build_route, 

151 wiki_update_route, 

152 wiki_status_route, 

153 wiki_synthesize_route, 

154 ], 

155 request_max_body_size=10 * 1024 * 1024, 

156 cors_config=cors, 

157 openapi_config=OpenAPIConfig( 

158 title="lilbee", 

159 description="Local knowledge base REST API", 

160 version=get_version(), 

161 path="/schema", 

162 ), 

163 )