Coverage for src/lilbee/server/routes/search.py: 100%
102 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Search, ask, ask_stream, chat, and chat_stream route handlers.
3Every route needs the token: these return the user's own documents.
4"""
6from __future__ import annotations
8import logging
9from collections.abc import AsyncGenerator
10from typing import Annotated, NoReturn
12from litestar import get, post
13from litestar.background_tasks import BackgroundTask
14from litestar.exceptions import HTTPException, ValidationException
15from litestar.params import FromQuery, QueryParameter
16from litestar.response import Stream
18from lilbee.core.results import DocumentResult
19from lilbee.data.store import EmbeddingModelMismatchError
20from lilbee.providers.base import ProviderError, ProviderErrorKind
21from lilbee.retrieval.query import ChatMessage as ChatMessageDict
22from lilbee.server import handlers
23from lilbee.server.chat_completions_api.errors import classify_provider_error
24from lilbee.server.chat_dispatch.concurrency import (
25 ChatBusyError,
26 ChatSlotGuard,
27 acquire_chat_slot_or_busy,
28 release_chat_slot,
29)
30from lilbee.server.chat_dispatch.dispatch import (
31 ModelDoesNotSupportToolsError,
32 ModelNotFoundError,
33)
34from lilbee.server.handlers.sse import SSE_MEDIA_TYPE, sse_error
35from lilbee.server.models import (
36 AskRequest,
37 AskResponse,
38 ChatRequest,
39 decode_chunk_type,
40)
42_BAD_REQUEST_STATUS = 400
43_NOT_FOUND_STATUS = 404
44_SERVICE_UNAVAILABLE_STATUS = 503
46# Shipped clients read /api 401/429 as lilbee-session signals, so upstream kinds stay 503.
47_API_PROVIDER_KIND_STATUSES: dict[ProviderErrorKind, int] = {
48 ProviderErrorKind.CONTEXT_OVERFLOW: _BAD_REQUEST_STATUS,
49 ProviderErrorKind.NOT_FOUND: _NOT_FOUND_STATUS,
50}
52log = logging.getLogger(__name__)
55def _embedding_mismatch_http(exc: EmbeddingModelMismatchError) -> HTTPException:
56 """Translate an embedder mismatch into a 409 carrying the facts to adopt.
58 The client renders its own confirm-to-adopt prompt from ``extra`` and, on
59 confirm, sets the embedder via ``PUT /api/models/embedding`` then retries.
60 The server never switches embedder unprompted.
61 """
62 return HTTPException(
63 status_code=409,
64 detail=str(exc),
65 extra={
66 "persisted_model": exc.persisted_model,
67 "persisted_dim": exc.persisted_dim,
68 "current_model": exc.current_model,
69 "adoptable": exc.dims_match,
70 },
71 )
74def _raise_chat_http_error(exc: Exception) -> NoReturn:
75 """Translate a chat/RAG failure into the Litestar HTTP envelope.
77 ValueError is a 422 validation error; a typed dispatch error or a
78 kind-mapped ProviderError carries its own status; anything else is a
79 503 carrying the failure message.
80 """
81 if isinstance(exc, ValueError):
82 raise ValidationException(str(exc)) from exc
83 raise HTTPException(status_code=_api_chat_error_status(exc), detail=str(exc)) from exc
86def _api_chat_error_status(exc: Exception) -> int:
87 """HTTP status for a non-stream /api chat failure; unmapped kinds stay 503."""
88 if isinstance(exc, ModelNotFoundError):
89 return _NOT_FOUND_STATUS
90 if isinstance(exc, ModelDoesNotSupportToolsError):
91 return _BAD_REQUEST_STATUS
92 if isinstance(exc, ProviderError):
93 return _API_PROVIDER_KIND_STATUSES.get(exc.kind, _SERVICE_UNAVAILABLE_STATUS)
94 return _SERVICE_UNAVAILABLE_STATUS
97async def _acquire_chat_lock_or_raise() -> None:
98 """Translate the canonical busy signal into Litestar's HTTP 429 envelope."""
99 from lilbee.app.services import get_services
101 try:
102 await acquire_chat_slot_or_busy(get_services().provider.max_concurrent_chats())
103 except ChatBusyError as exc:
104 raise HTTPException(status_code=429, detail=str(exc), headers={"Retry-After": "1"}) from exc
107async def _gated_stream(
108 generator: AsyncGenerator[str, None],
109 guard: ChatSlotGuard,
110) -> AsyncGenerator[str, None]:
111 """Wrap *generator* so the chat lock is released when the stream ends.
113 The lock must already be held when this is called. Release happens on
114 natural completion, exception, and client-disconnect (GeneratorExit
115 fires the ``finally`` block); a disconnect before the first iteration
116 never enters this body, so the route also releases *guard* from the
117 response's after-send hook. A failure inside the generator becomes an
118 SSE error event; raising after the 201 headers would drop the connection
119 with no body for the client to read.
120 """
121 try:
122 async for chunk in generator:
123 yield chunk
124 except Exception as exc:
125 log.exception("streaming chat handler failed")
126 # Same mapper as the non-streaming sibling: without a code the stream
127 # flattened unknown-model, no-tool-support and context-overflow into
128 # one message, so a client could not tell "pull the model" from
129 # "shorten the prompt". It also redacts the backend-failure kinds.
130 classified = classify_provider_error(exc)
131 if classified is None:
132 yield sse_error(str(exc))
133 else:
134 yield sse_error(classified.message, code=classified.code)
135 finally:
136 await guard.release()
139def _slot_gated_sse(generator: AsyncGenerator[str, None], guard: ChatSlotGuard) -> Stream:
140 """SSE Stream whose chat slot is freed by the generator or the after-send hook."""
141 return Stream(
142 _gated_stream(generator, guard),
143 media_type=SSE_MEDIA_TYPE,
144 background=BackgroundTask(guard.release),
145 )
148@get("/api/search")
149async def search_route(
150 q: FromQuery[str],
151 top_k: Annotated[int, QueryParameter(ge=1, le=100)] = 5,
152 chunk_type: FromQuery[str | None] = None,
153) -> list[DocumentResult]:
154 """Search indexed documents by semantic similarity. No LLM call required."""
155 try:
156 parsed_chunk_type = decode_chunk_type(chunk_type)
157 except ValueError as exc:
158 raise ValidationException(str(exc)) from exc
159 try:
160 return await handlers.search(q, top_k=top_k, chunk_type=parsed_chunk_type)
161 except EmbeddingModelMismatchError as exc:
162 # ``adoptable`` is a bare yes/no on whether switching embedder alone
163 # fixes the index, which is what a client renders its hint from. The
164 # embedder names stay out of the generic 503.
165 raise HTTPException(
166 status_code=409,
167 detail="The index was built with a different embedding model than the one configured.",
168 extra={"adoptable": exc.dims_match},
169 ) from exc
170 except ValueError as exc:
171 raise ValidationException(str(exc)) from exc
172 except Exception as exc:
173 # str(exc) here routinely carries data-root paths, LanceDB table names,
174 # and model ids. Log the real cause for the operator; return a generic
175 # message on the wire.
176 log.exception("Search failed")
177 raise HTTPException(status_code=503, detail="Search is temporarily unavailable.") from exc
180@post("/api/ask")
181async def ask_route(data: AskRequest) -> AskResponse:
182 """One-shot RAG question returning an answer with source chunks."""
183 await _acquire_chat_lock_or_raise()
184 try:
185 return await handlers.ask(
186 question=data.question,
187 top_k=data.top_k,
188 options=data.options,
189 chunk_type=data.chunk_type,
190 )
191 except EmbeddingModelMismatchError as exc:
192 raise _embedding_mismatch_http(exc) from exc
193 except Exception as exc:
194 # No separate ValueError arm: _raise_chat_http_error is the single
195 # translation point and its first branch is ValueError -> 422, which is
196 # what chat_route relies on. Two copies drift.
197 _raise_chat_http_error(exc)
198 finally:
199 await release_chat_slot()
202@post("/api/ask/stream", media_type=SSE_MEDIA_TYPE)
203async def ask_stream_route(data: AskRequest) -> Stream:
204 """Streaming SSE version of ask, emitting token-by-token answer chunks."""
205 await _acquire_chat_lock_or_raise()
206 return _slot_gated_sse(
207 handlers.ask_stream(
208 question=data.question,
209 top_k=data.top_k,
210 options=data.options,
211 chunk_type=data.chunk_type,
212 ),
213 ChatSlotGuard(),
214 )
217@post("/api/chat")
218async def chat_route(data: ChatRequest) -> AskResponse:
219 """RAG chat with conversation history, returning an answer with sources."""
220 await _acquire_chat_lock_or_raise()
221 history: list[ChatMessageDict] = [
222 ChatMessageDict(role=m.role, content=m.content) for m in data.history
223 ]
224 try:
225 return await handlers.chat(
226 question=data.question,
227 history=history,
228 top_k=data.top_k,
229 options=data.options,
230 chunk_type=data.chunk_type,
231 summary=data.summary,
232 session_id=data.session_id,
233 )
234 except EmbeddingModelMismatchError as exc:
235 raise _embedding_mismatch_http(exc) from exc
236 except Exception as exc:
237 _raise_chat_http_error(exc)
238 finally:
239 await release_chat_slot()
242@post("/api/chat/stream", media_type=SSE_MEDIA_TYPE)
243async def chat_stream_route(data: ChatRequest) -> Stream:
244 """Streaming SSE version of chat with conversation history."""
245 await _acquire_chat_lock_or_raise()
246 history: list[ChatMessageDict] = [
247 ChatMessageDict(role=m.role, content=m.content) for m in data.history
248 ]
249 return _slot_gated_sse(
250 handlers.chat_stream(
251 question=data.question,
252 history=history,
253 top_k=data.top_k,
254 options=data.options,
255 chunk_type=data.chunk_type,
256 summary=data.summary,
257 session_id=data.session_id,
258 ),
259 ChatSlotGuard(),
260 )