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

1"""Document management route handlers: add, list, remove, sync, export. 

2 

3Every route needs the token, reads included: the listing names the user's 

4files and ``/api/export`` serializes the whole corpus. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from typing import Annotated 

11 

12from litestar import Request, Response, get, post 

13from litestar.datastructures import UploadFile 

14from litestar.exceptions import ValidationException 

15from litestar.params import FromQuery, MultipartBody, QueryParameter 

16from litestar.response import Stream 

17from pydantic import BaseModel, Field 

18 

19from lilbee.server import handlers 

20from lilbee.server.handlers.sse import SSE_MEDIA_TYPE 

21from lilbee.server.models import ( 

22 AddRequest, 

23 DocumentListResponse, 

24 DocumentRemoveResponse, 

25 SyncRequest, 

26) 

27 

28 

29class RemoveRequest(BaseModel): 

30 """Request body for /api/documents/remove.""" 

31 

32 names: list[str] = Field(max_length=100) 

33 

34 

35@post("/api/sync", media_type=SSE_MEDIA_TYPE) 

36async def sync_route(data: SyncRequest | None = None) -> Stream: 

37 """Re-index changed documents with streaming SSE progress events. 

38 

39 Pass ``{"force_rebuild": true}`` to wipe the store and re-ingest every file 

40 under the current ``cfg.embedding_model``. This is the recovery path after 

41 a ``PUT /api/models/embedding`` that returned ``reindex_required=true``. 

42 Pass ``{"retry_skipped": true}`` for the lighter path: retry the files that 

43 failed a previous sync without dropping the store. 

44 Pass ``{"prune_ignored": true}`` to also drop sources a ``.lilbeeignore`` 

45 now excludes; without it, sync leaves already-indexed sources alone. 

46 """ 

47 enable_ocr = data.enable_ocr if data else None 

48 force_rebuild = data.force_rebuild if data else False 

49 retry_skipped = data.retry_skipped if data else False 

50 prune_ignored = data.prune_ignored if data else False 

51 return Stream( 

52 handlers.sync_stream( 

53 enable_ocr=enable_ocr, 

54 force_rebuild=force_rebuild, 

55 retry_skipped=retry_skipped, 

56 prune_ignored=prune_ignored, 

57 ), 

58 media_type=SSE_MEDIA_TYPE, 

59 ) 

60 

61 

62@post("/api/add", media_type=SSE_MEDIA_TYPE) 

63async def add_route(data: AddRequest) -> Stream: 

64 """Add files to the knowledge base with streaming SSE progress.""" 

65 try: 

66 paths, force, enable_ocr, ocr_timeout = handlers.validate_add_paths(data.model_dump()) 

67 except ValueError as exc: 

68 raise ValidationException(str(exc)) from exc 

69 return Stream( 

70 handlers.add_files_stream( 

71 paths, force=force, enable_ocr=enable_ocr, ocr_timeout=ocr_timeout 

72 ), 

73 media_type=SSE_MEDIA_TYPE, 

74 status_code=201, 

75 ) 

76 

77 

78@post("/api/add/upload", media_type=SSE_MEDIA_TYPE) 

79async def add_upload_route( 

80 data: MultipartBody[list[UploadFile]], 

81) -> Stream: 

82 """Ingest uploaded file content with streaming SSE progress. 

83 

84 Unlike /api/add, which reads server-side paths, this accepts the client's raw 

85 file bytes. That lets a client whose files the server cannot read by path -- 

86 e.g. the plugin or CLI in external mode against a remote lilbee / GPU box -- 

87 ingest its own local files by uploading them straight to the server. 

88 """ 

89 # Names first, bytes second: reading every part before validating cost a 

90 # full in-memory copy of a payload that was going to be rejected anyway. 

91 try: 

92 names = handlers.validate_upload_names([upload.filename for upload in data]) 

93 except ValueError as exc: 

94 raise ValidationException(str(exc)) from exc 

95 cleaned = [(name, await upload.read()) for name, upload in zip(names, data, strict=True)] 

96 return Stream( 

97 handlers.add_uploads_stream(cleaned), 

98 media_type=SSE_MEDIA_TYPE, 

99 status_code=201, 

100 ) 

101 

102 

103@get("/api/documents") 

104async def documents_list_route( 

105 search: FromQuery[str] = "", 

106 limit: Annotated[int, QueryParameter(ge=1, le=1000)] = 50, 

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

108) -> DocumentListResponse: 

109 """List indexed documents with metadata, paginated and searchable.""" 

110 return await handlers.list_documents(search=search, limit=limit, offset=offset) 

111 

112 

113@post("/api/documents/remove") 

114async def documents_remove_route(data: RemoveRequest) -> DocumentRemoveResponse: 

115 """Remove documents from the knowledge base by source name.""" 

116 return await handlers.delete_documents(data.names) 

117 

118 

119@get("/api/export", media_type="application/octet-stream") 

120async def export_route( 

121 fmt: Annotated[str, QueryParameter(name="format")] = "", 

122 source: FromQuery[str] = "", 

123) -> Response[bytes]: 

124 """Download the per-page text dataset as a file (parquet by default). 

125 

126 The media type is declared on the decorator as well as on the returned 

127 Response for the same reason the streaming routes declare theirs: litestar 

128 documents the content type from the decorator, so without it the schema 

129 promises JSON and a generated client parses a parquet file as text. 

130 """ 

131 from lilbee.app.dataset import DatasetError, export_to_bytes 

132 

133 try: 

134 # export_to_bytes serializes the whole per-page dataset into memory; 

135 # offload so a large export doesn't stall every other request, matching 

136 # get_source_content's own off-loop read. 

137 payload = await asyncio.to_thread(export_to_bytes, fmt, source or None) 

138 except DatasetError as exc: 

139 raise ValidationException(str(exc)) from exc 

140 return Response( 

141 content=payload.data, 

142 media_type="application/octet-stream", 

143 headers={"content-disposition": f'attachment; filename="pages.{payload.fmt}"'}, 

144 ) 

145 

146 

147@post("/api/import", media_type=SSE_MEDIA_TYPE) 

148async def import_route( 

149 request: Request, 

150 fmt: Annotated[str, QueryParameter(name="format")] = "", 

151) -> Stream: 

152 """Import an uploaded per-page dataset with streaming SSE progress events. 

153 

154 The request body is the raw dataset bytes; ``?format=parquet|jsonl`` is 

155 required since there is no filename to infer from. Bounded by the server's 

156 body-size limit; larger datasets use the path-based CLI/MCP import. 

157 """ 

158 from lilbee.app.dataset import DatasetError, require_format 

159 

160 try: 

161 require_format(fmt) 

162 except DatasetError as exc: 

163 raise ValidationException(str(exc)) from exc 

164 return Stream( 

165 handlers.import_stream(await request.body(), fmt), 

166 media_type=SSE_MEDIA_TYPE, 

167 status_code=201, 

168 )