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 """ 

45 enable_ocr = data.enable_ocr if data else None 

46 force_rebuild = data.force_rebuild if data else False 

47 retry_skipped = data.retry_skipped if data else False 

48 return Stream( 

49 handlers.sync_stream( 

50 enable_ocr=enable_ocr, force_rebuild=force_rebuild, retry_skipped=retry_skipped 

51 ), 

52 media_type=SSE_MEDIA_TYPE, 

53 ) 

54 

55 

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

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

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

59 try: 

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

61 except ValueError as exc: 

62 raise ValidationException(str(exc)) from exc 

63 return Stream( 

64 handlers.add_files_stream( 

65 paths, force=force, enable_ocr=enable_ocr, ocr_timeout=ocr_timeout 

66 ), 

67 media_type=SSE_MEDIA_TYPE, 

68 status_code=201, 

69 ) 

70 

71 

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

73async def add_upload_route( 

74 data: MultipartBody[list[UploadFile]], 

75) -> Stream: 

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

77 

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

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

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

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

82 """ 

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

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

85 try: 

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

87 except ValueError as exc: 

88 raise ValidationException(str(exc)) from exc 

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

90 return Stream( 

91 handlers.add_uploads_stream(cleaned), 

92 media_type=SSE_MEDIA_TYPE, 

93 status_code=201, 

94 ) 

95 

96 

97@get("/api/documents") 

98async def documents_list_route( 

99 search: FromQuery[str] = "", 

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

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

102) -> DocumentListResponse: 

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

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

105 

106 

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

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

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

110 return await handlers.delete_documents(data.names) 

111 

112 

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

114async def export_route( 

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

116 source: FromQuery[str] = "", 

117) -> Response[bytes]: 

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

119 

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

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

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

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

124 """ 

125 from lilbee.app.dataset import DatasetError, export_to_bytes 

126 

127 try: 

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

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

130 # get_source_content's own off-loop read. 

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

132 except DatasetError as exc: 

133 raise ValidationException(str(exc)) from exc 

134 return Response( 

135 content=payload.data, 

136 media_type="application/octet-stream", 

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

138 ) 

139 

140 

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

142async def import_route( 

143 request: Request, 

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

145) -> Stream: 

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

147 

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

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

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

151 """ 

152 from lilbee.app.dataset import DatasetError, require_format 

153 

154 try: 

155 require_format(fmt) 

156 except DatasetError as exc: 

157 raise ValidationException(str(exc)) from exc 

158 return Stream( 

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

160 media_type=SSE_MEDIA_TYPE, 

161 status_code=201, 

162 )