diff --git a/conf/config.yaml b/conf/config.yaml index 0cd663eaa..5567f4e27 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -14,9 +14,12 @@ _llm_params: &llm_params timeout: 60 max_retries: 2 logprobs: true + # null omits chat_template_kwargs; set false for Qwen-style models when + # reasoning traces should be disabled. Leave unset for Mistral tokenizers. + enable_thinking: null # --- LLM --- -# Env: BASE_URL, MODEL, API_KEY +# Env: BASE_URL, MODEL, API_KEY, LLM_ENABLE_THINKING llm: <<: *llm_params base_url: "" @@ -24,7 +27,7 @@ llm: api_key: "" # --- VLM (Vision Language Model) --- -# Env: VLM_BASE_URL, VLM_MODEL, VLM_API_KEY +# Env: VLM_BASE_URL, VLM_MODEL, VLM_API_KEY, VLM_ENABLE_THINKING vlm: <<: *llm_params base_url: "" @@ -239,7 +242,8 @@ loader: # Env: OPENAI_LOADER_BASE_URL, OPENAI_LOADER_API_KEY, OPENAI_LOADER_MODEL, # OPENAI_LOADER_TEMPERATURE, OPENAI_LOADER_TIMEOUT, OPENAI_LOADER_MAX_RETRIES, - # OPENAI_LOADER_TOP_P, OPENAI_LOADER_CONCURRENCY_LIMIT + # OPENAI_LOADER_TOP_P, OPENAI_LOADER_CONCURRENCY_LIMIT, + # OPENAI_LOADER_ENABLE_THINKING openai: base_url: http://openai:8000/v1 api_key: EMPTY @@ -249,6 +253,7 @@ loader: max_retries: 2 top_p: 0.9 concurrency_limit: 20 + enable_thinking: null # --- Ray --- ray: diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index e55fbb444..de5c1527d 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -68,6 +68,7 @@ The parameters below configure how the OCR loader communicates with the model se | `OPENAI_LOADER_MAX_RETRIES` | `int` | `2` | Number of retry attempts for failed OCR requests. | | `OPENAI_LOADER_TOP_P` | `float` | `0.9` | Nucleus sampling parameter that limits generation to the top-p probability mass. | | `OPENAI_LOADER_CONCURRENCY_LIMIT` | `int` | `20` | Maximum number of OCR requests processed concurrently. Useful for multi-page PDF workloads. | +| `OPENAI_LOADER_ENABLE_THINKING` | `bool` | unset | Optional chat-template control for OCR VLM models that support `enable_thinking`; leave unset for Mistral tokenizers, set `false` to suppress Qwen-style reasoning traces. | :::note[Information] This feature is currently experimental. Docker server configurations are available in [extern/ocr_vlm_servers](https://github.com/linagora/openrag/tree/main/extern/ocr_vlm_servers) and can be deployed using standard Docker Compose commands. @@ -235,6 +236,7 @@ These are external services to provide !!! | `BASE_URL` | str | Base URL of the LLM API endpoint | | `MODEL` | str | Model identifier for the LLM | | `API_KEY` | str | API key for authenticating with the LLM service | +| `LLM_ENABLE_THINKING` | bool | Optional chat-template control for models that support `enable_thinking`; leave unset for Mistral tokenizers, set `false` to suppress Qwen-style reasoning traces | | `LLM_SEMAPHORE` | int | 10 | Maximum number of concurrent requests to allow for the LLM service | | `MAX_LLM_CONTEXT_SIZE` | `int` | `8192` | Fallback maximum token limit for chat/completion requests. At startup, the `/v1/models` endpoint is queried for the model's `max_model_len`; if that query fails this value is used instead. Requests whose total token count (prompt + `max_tokens`) exceeds the limit are rejected with a **413** error. | @@ -245,6 +247,7 @@ These are external services to provide !!! | `VLM_BASE_URL` | str | Base URL of the VLM API endpoint | | `VLM_MODEL` | str | Model identifier for the VLM | | `VLM_API_KEY` | str | API key for authenticating with the VLM service | +| `VLM_ENABLE_THINKING` | bool | Optional chat-template control for models that support `enable_thinking`; leave unset for Mistral tokenizers, set `false` to suppress Qwen-style reasoning traces | | `VLM_SEMAPHORE` | int | 10 | Maximum number of concurrent requests to allow for the VLM service | ### Retriever Configuration diff --git a/infra/compose/.env.example b/infra/compose/.env.example index c50fc7b26..5ee3429a8 100644 --- a/infra/compose/.env.example +++ b/infra/compose/.env.example @@ -2,11 +2,20 @@ BASE_URL= API_KEY= MODEL= +# Optional: set false for Qwen-style models to suppress reasoning traces. +# Leave unset for Mistral tokenizers. +# LLM_ENABLE_THINKING=false # VLM (Visual Language Model) you can set it to the same as LLM if your LLM supports images VLM_BASE_URL= VLM_API_KEY= VLM_MODEL= +# Optional: same behavior as LLM_ENABLE_THINKING for VLM chat templates. +# VLM_ENABLE_THINKING=false + +# OCR VLM loader (DotsOCR/OpenAILoader) thinking control. +# Leave unset for Mistral tokenizers. +# OPENAI_LOADER_ENABLE_THINKING=false ## FastAPI App (no need to change it) # APP_PORT=8080 # this is the forwarded port diff --git a/openrag/core/config/endpoints.py b/openrag/core/config/endpoints.py index 25ad2e19a..5e521096a 100644 --- a/openrag/core/config/endpoints.py +++ b/openrag/core/config/endpoints.py @@ -14,6 +14,7 @@ class LLMParamsConfig(ConfigMixin): timeout: int = 60 max_retries: int = 2 logprobs: bool = True + enable_thinking: bool | None = None class LLMConfig(LLMParamsConfig): diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py index 5db92d298..0469ca5ea 100644 --- a/openrag/core/config/indexation.py +++ b/openrag/core/config/indexation.py @@ -58,6 +58,7 @@ class OpenAILoaderConfig(ConfigMixin): max_retries: int = 2 top_p: float = 0.9 concurrency_limit: int = 20 + enable_thinking: bool | None = None # --------------------------------------------------------------------------- diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py index 5e0c5900a..c5d8beb61 100644 --- a/openrag/core/config/loader.py +++ b/openrag/core/config/loader.py @@ -27,10 +27,12 @@ ("BASE_URL", "llm.base_url", str), ("MODEL", "llm.model", str), ("API_KEY", "llm.api_key", str), + ("LLM_ENABLE_THINKING", "llm.enable_thinking", bool), # VLM ("VLM_BASE_URL", "vlm.base_url", str), ("VLM_MODEL", "vlm.model", str), ("VLM_API_KEY", "vlm.api_key", str), + ("VLM_ENABLE_THINKING", "vlm.enable_thinking", bool), # Semaphore ("LLM_SEMAPHORE", "semaphore.llm_semaphore", int), ("VLM_SEMAPHORE", "semaphore.vlm_semaphore", int), @@ -123,6 +125,7 @@ ("OPENAI_LOADER_MAX_RETRIES", "loader.openai.max_retries", int), ("OPENAI_LOADER_TOP_P", "loader.openai.top_p", float), ("OPENAI_LOADER_CONCURRENCY_LIMIT", "loader.openai.concurrency_limit", int), + ("OPENAI_LOADER_ENABLE_THINKING", "loader.openai.enable_thinking", bool), # Ray ("RAY_NUM_GPUS", "ray.num_gpus", float), ("RAY_POOL_SIZE", "ray.pool_size", int), diff --git a/openrag/core/utils/text.py b/openrag/core/utils/text.py index 656d39f68..1c499363e 100644 --- a/openrag/core/utils/text.py +++ b/openrag/core/utils/text.py @@ -28,7 +28,9 @@ def get_num_tokens(): from langchain_openai import ChatOpenAI config = load_config() - llm = ChatOpenAI(**config.llm.model_dump()) + llm_kwargs = config.llm.model_dump() + llm_kwargs.pop("enable_thinking", None) + llm = ChatOpenAI(**llm_kwargs) _cached_length_function = llm.get_num_tokens except Exception as exc: import tiktoken diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py index a01764130..0f331a161 100644 --- a/openrag/services/inference/vllm_client.py +++ b/openrag/services/inference/vllm_client.py @@ -64,11 +64,13 @@ def __init__( *, api_key: str = "", timeout: float = 240.0, + enable_thinking: bool | None = None, **kwargs, ) -> None: self._endpoint = endpoint.rstrip("/") self._model = model_name self._api_key = api_key + self._enable_thinking = enable_thinking self._defaults: dict = kwargs headers: dict[str, str] = {"Content-Type": "application/json"} if api_key: @@ -101,6 +103,15 @@ def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str] | N return base_url, model, override_headers + def _chat_payload_kwargs(self, kwargs: dict) -> dict: + payload_kwargs = {**self._defaults, **kwargs} + enable_thinking = payload_kwargs.pop("enable_thinking", self._enable_thinking) + if enable_thinking is not None: + chat_template_kwargs = dict(payload_kwargs.get("chat_template_kwargs") or {}) + chat_template_kwargs.setdefault("enable_thinking", enable_thinking) + payload_kwargs["chat_template_kwargs"] = chat_template_kwargs + return payload_kwargs + @with_circuit_breaker("llm") @with_retry(max_attempts=3) async def generate(self, prompt: str, **kwargs) -> dict: @@ -126,7 +137,7 @@ async def generate(self, prompt: str, **kwargs) -> dict: async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: base_url, model, headers = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) - payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": False} + payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": False} try: resp = await self._client.post(f"{base_url}/chat/completions", json=payload, headers=headers) resp.raise_for_status() @@ -144,7 +155,7 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: base_url, model, headers = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) - payload = {**self._defaults, **kwargs, "model": model, "messages": messages, "stream": True} + payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": True} try: async with self._client.stream( "POST", f"{base_url}/chat/completions", json=payload, headers=headers @@ -357,10 +368,16 @@ async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> ], } ] + payload: dict = { + **self._chat_payload_kwargs({}), + "model": self._model, + "messages": messages, + "max_tokens": self._max_tokens, + } try: resp = await self._client.post( f"{self._endpoint}/chat/completions", - json={"model": self._model, "messages": messages, "max_tokens": self._max_tokens}, + json=payload, ) resp.raise_for_status() except httpx.ConnectError as exc: diff --git a/openrag/services/orchestrators/model_endpoint_service.py b/openrag/services/orchestrators/model_endpoint_service.py index 64dc34a4b..6a27098b7 100644 --- a/openrag/services/orchestrators/model_endpoint_service.py +++ b/openrag/services/orchestrators/model_endpoint_service.py @@ -37,6 +37,13 @@ def _with_api_key(extra: dict[str, Any], api_key: str | None) -> dict[str, Any]: return extra +def _with_enable_thinking(extra: dict[str, Any], enable_thinking: bool | None) -> dict[str, Any]: + """Add chat-template thinking control only when explicitly configured.""" + if enable_thinking is None: + return extra + return {**extra, "enable_thinking": enable_thinking} + + class ModelEndpointService: """CRUD and lifecycle management for named model endpoints.""" @@ -103,17 +110,23 @@ def _build_default_seeds(self) -> dict[str, dict[str, Any]]: "llm": { "endpoint": os.getenv("LLM_ENDPOINT", s.llm.base_url), "model_name": os.getenv("LLM_MODEL", s.llm.model), - "extra": _with_api_key( - {"implementation": "vllm"}, - os.getenv("API_KEY", s.llm.api_key), + "extra": _with_enable_thinking( + _with_api_key( + {"implementation": "vllm"}, + os.getenv("API_KEY", s.llm.api_key), + ), + s.llm.enable_thinking, ), }, "vlm": { "endpoint": os.getenv("VLM_ENDPOINT", s.vlm.base_url), "model_name": os.getenv("VLM_MODEL", s.vlm.model), - "extra": _with_api_key( - {"implementation": "vllm"}, - os.getenv("VLM_API_KEY", s.vlm.api_key), + "extra": _with_enable_thinking( + _with_api_key( + {"implementation": "vllm"}, + os.getenv("VLM_API_KEY", s.vlm.api_key), + ), + s.vlm.enable_thinking, ), }, "reranker": { diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index 486fafb01..9f448b073 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -308,6 +308,9 @@ def _global_llm_endpoint_config(cfg: Any) -> Any | None: "max_retries": getattr(llm_cfg, "max_retries", 2), "logprobs": getattr(llm_cfg, "logprobs", True), } + enable_thinking = getattr(llm_cfg, "enable_thinking", None) + if enable_thinking is not None: + extra["enable_thinking"] = enable_thinking return ModelEndpointConfig( endpoint=endpoint, model_name=model_name, diff --git a/openrag/services/workers/parsers/parser_dispatcher.py b/openrag/services/workers/parsers/parser_dispatcher.py index d79ea076f..713c4bd63 100644 --- a/openrag/services/workers/parsers/parser_dispatcher.py +++ b/openrag/services/workers/parsers/parser_dispatcher.py @@ -169,7 +169,13 @@ def _build_pdf_client(self) -> DocumentParser: from services.inference.parsers.dotsocr import DotsOCRPdfClient ocfg = self._config.loader.openai - vlm = _build_vlm(ocfg.base_url, ocfg.model, ocfg.api_key, ocfg.timeout) + vlm = _build_vlm( + ocfg.base_url, + ocfg.model, + ocfg.api_key, + ocfg.timeout, + ocfg.enable_thinking, + ) client = DotsOCRPdfClient(vlm, concurrency_limit=ocfg.concurrency_limit) return _create("core.indexing.parsers.pdf.client_based", "pdf_client", client=client) @@ -215,12 +221,19 @@ def _suffix(filename: str) -> str: return filename.rsplit(".", 1)[-1].lower() if "." in filename else "" -def _build_vlm(base_url: str, model: str, api_key: str, timeout: float) -> Any: +def _build_vlm(base_url: str, model: str, api_key: str, timeout: float, enable_thinking: bool | None = None) -> Any: """Construct a vLLM-backed VLM client for VLM-OCR / captioning.""" import services.inference.vllm_client # noqa: F401 - registers "vllm" from core.vlm import vlm_registry - return vlm_registry.create("vllm", endpoint=base_url, model_name=model, api_key=api_key, timeout=timeout) + return vlm_registry.create( + "vllm", + endpoint=base_url, + model_name=model, + api_key=api_key, + timeout=timeout, + enable_thinking=enable_thinking, + ) def build_parser_dispatcher(config: Any) -> ParserDispatcher: @@ -244,7 +257,13 @@ def build_caption_vlm(config: Any) -> Any | None: vlm_cfg = config.vlm if not getattr(vlm_cfg, "base_url", ""): return None - return _build_vlm(vlm_cfg.base_url, vlm_cfg.model, vlm_cfg.api_key, vlm_cfg.timeout) + return _build_vlm( + vlm_cfg.base_url, + vlm_cfg.model, + vlm_cfg.api_key, + vlm_cfg.timeout, + vlm_cfg.enable_thinking, + ) __all__ = ["ParserDispatcher", "build_parser_dispatcher", "build_caption_vlm"] diff --git a/tests/unit/core/config/test_rdb_env.py b/tests/unit/core/config/test_rdb_env.py index 54766225e..61f4dc7fc 100644 --- a/tests/unit/core/config/test_rdb_env.py +++ b/tests/unit/core/config/test_rdb_env.py @@ -12,3 +12,16 @@ def test_postgres_provisioning_flags_can_be_overridden_from_env(monkeypatch, tmp assert settings.rdb.auto_create_database is False assert settings.rdb.run_migrations is False + + +def test_llm_and_vlm_enable_thinking_can_be_overridden_from_env(monkeypatch, tmp_path): + (tmp_path / "config.yaml").write_text("retriever:\n type: single\n", encoding="utf-8") + monkeypatch.setenv("LLM_ENABLE_THINKING", "false") + monkeypatch.setenv("VLM_ENABLE_THINKING", "true") + monkeypatch.setenv("OPENAI_LOADER_ENABLE_THINKING", "false") + + settings = load_config(config_path=tmp_path) + + assert settings.llm.enable_thinking is False + assert settings.vlm.enable_thinking is True + assert settings.loader.openai.enable_thinking is False diff --git a/tests/unit/core/utils/test_text.py b/tests/unit/core/utils/test_text.py new file mode 100644 index 000000000..29b0e62ee --- /dev/null +++ b/tests/unit/core/utils/test_text.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import sys +from types import SimpleNamespace + + +def test_get_num_tokens_does_not_pass_enable_thinking_to_chatopenai(monkeypatch): + import core.utils.text as text_utils + + captured: dict = {} + + class FakeChatOpenAI: + def __init__(self, **kwargs): + captured.update(kwargs) + + def get_num_tokens(self, _text: str) -> int: + return 1 + + monkeypatch.setitem(sys.modules, "langchain_openai", SimpleNamespace(ChatOpenAI=FakeChatOpenAI)) + monkeypatch.setattr( + text_utils, + "load_config", + lambda: SimpleNamespace( + llm=SimpleNamespace( + model_dump=lambda: { + "base_url": "http://llm:8000/v1", + "model": "qwen", + "api_key": "key", + "enable_thinking": False, + } + ) + ), + ) + monkeypatch.setattr(text_utils, "_cached_length_function", None) + + length_function = text_utils.get_num_tokens() + + assert length_function("hello") == 1 + assert "enable_thinking" not in captured diff --git a/tests/unit/services/inference/test_vllm_client.py b/tests/unit/services/inference/test_vllm_client.py index 39b038957..913a25320 100644 --- a/tests/unit/services/inference/test_vllm_client.py +++ b/tests/unit/services/inference/test_vllm_client.py @@ -102,6 +102,20 @@ def handler(request: httpx.Request) -> httpx.Response: assert 'data: {"choices":[{"delta":{"content":"Hello"}}]}' in lines assert 'data: {"choices":[{"delta":{"content":" world"}}]}' in lines + @pytest.mark.asyncio + async def test_stream_chat_sends_enable_thinking_as_chat_template_kwargs_when_configured(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["stream"] is True + assert body["chat_template_kwargs"] == {"enable_thinking": False} + assert "enable_thinking" not in body + return httpx.Response(200, text="data: [DONE]\n") + + client = self._make_client(handler, enable_thinking=False) + lines = [line async for line in client.stream_chat([{"role": "user", "content": "hi"}])] + + assert lines == ["data: [DONE]"] + @pytest.mark.asyncio async def test_stream_chat_error_raises(self): client = self._make_client(lambda req: httpx.Response(503, text="unavailable")) @@ -142,6 +156,47 @@ def capture(req: httpx.Request) -> httpx.Response: await self._make_client(capture).chat([{"role": "user", "content": "hi"}]) assert captured["temperature"] == 0.3 + @pytest.mark.asyncio + async def test_enable_thinking_is_sent_as_chat_template_kwargs_when_configured(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client(capture, enable_thinking=False).chat([{"role": "user", "content": "hi"}]) + + assert captured["chat_template_kwargs"] == {"enable_thinking": False} + assert "enable_thinking" not in captured + + @pytest.mark.asyncio + async def test_enable_thinking_merges_with_existing_chat_template_kwargs(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client( + capture, + enable_thinking=False, + chat_template_kwargs={"custom": "value"}, + ).chat([{"role": "user", "content": "hi"}]) + + assert captured["chat_template_kwargs"] == {"custom": "value", "enable_thinking": False} + + @pytest.mark.asyncio + async def test_chat_template_kwargs_omitted_by_default(self): + captured: dict = {} + + def capture(req: httpx.Request) -> httpx.Response: + captured.update(json.loads(req.content)) + return _chat_response() + + await self._make_client(capture).chat([{"role": "user", "content": "hi"}]) + + assert "chat_template_kwargs" not in captured + @pytest.mark.asyncio async def test_per_request_kwargs_override_defaults(self): def handler(request: httpx.Request) -> httpx.Response: @@ -461,6 +516,37 @@ def handler(request: httpx.Request) -> httpx.Response: await self._make_vision(handler, max_tokens=512).caption_image(b"img") + @pytest.mark.asyncio + async def test_caption_image_sends_enable_thinking_as_chat_template_kwargs_when_configured(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["chat_template_kwargs"] == {"enable_thinking": False} + assert "enable_thinking" not in body + return _chat_response("ok") + + await self._make_vision(handler, enable_thinking=False).caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_image_merges_enable_thinking_with_existing_chat_template_kwargs(self): + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + assert body["chat_template_kwargs"] == {"custom": "value", "enable_thinking": False} + return _chat_response("ok") + + await self._make_vision( + handler, + enable_thinking=False, + chat_template_kwargs={"custom": "value"}, + ).caption_image(b"img") + + @pytest.mark.asyncio + async def test_caption_image_omits_chat_template_kwargs_by_default(self): + def handler(request: httpx.Request) -> httpx.Response: + assert "chat_template_kwargs" not in json.loads(request.content) + return _chat_response("ok") + + await self._make_vision(handler).caption_image(b"img") + @pytest.mark.asyncio async def test_caption_connection_error(self): async def fail(*a, **kw): diff --git a/tests/unit/services/orchestrators/test_model_endpoint_service.py b/tests/unit/services/orchestrators/test_model_endpoint_service.py index 548aafb11..4128a3cd2 100644 --- a/tests/unit/services/orchestrators/test_model_endpoint_service.py +++ b/tests/unit/services/orchestrators/test_model_endpoint_service.py @@ -180,6 +180,35 @@ async def test_seed_defaults_preserves_endpoint_api_keys(monkeypatch): } +@pytest.mark.asyncio +async def test_seed_defaults_preserves_llm_and_vlm_enable_thinking(monkeypatch): + from core.config.root import Settings + + monkeypatch.delenv("LLM_ENDPOINT", raising=False) + monkeypatch.delenv("LLM_MODEL", raising=False) + + settings = Settings( + llm={ + "base_url": "http://llm:8000/v1", + "model": "qwen", + "enable_thinking": False, + }, + vlm={ + "base_url": "http://vlm:8000/v1", + "model": "qwen-vl", + "enable_thinking": True, + }, + ) + repo = _FakeEndpointRepo() + svc = _make_service(repo, settings=settings) + + await svc.seed_defaults() + + rows = {row.model_type: row for row in repo._store.values()} + assert rows["llm"].extra["enable_thinking"] is False + assert rows["vlm"].extra["enable_thinking"] is True + + @pytest.mark.asyncio async def test_seed_defaults_skips_reranker_when_unconfigured(monkeypatch): from core.config.root import Settings diff --git a/tests/unit/services/workers/parsers/test_parser_dispatcher.py b/tests/unit/services/workers/parsers/test_parser_dispatcher.py index 1e0be51ac..802bcfcf0 100644 --- a/tests/unit/services/workers/parsers/test_parser_dispatcher.py +++ b/tests/unit/services/workers/parsers/test_parser_dispatcher.py @@ -13,7 +13,13 @@ def _config( - *, pdf="MarkerLoader", audio="LocalWhisperLoader", image_captioning=True, vlm_base_url="" + *, + pdf="MarkerLoader", + audio="LocalWhisperLoader", + image_captioning=True, + vlm_base_url="", + vlm_enable_thinking=None, + openai_loader_enable_thinking=None, ) -> SimpleNamespace: file_loaders = SimpleNamespace( pdf=pdf, @@ -22,8 +28,22 @@ def _config( flac=audio, mp4=audio, ) - loader = SimpleNamespace(file_loaders=file_loaders, image_captioning=image_captioning) - vlm = SimpleNamespace(base_url=vlm_base_url, model="m", api_key="k", timeout=60) + openai = SimpleNamespace( + base_url="http://openai:8000/v1", + model="dotsocr-model", + api_key="k", + timeout=60, + concurrency_limit=20, + enable_thinking=openai_loader_enable_thinking, + ) + loader = SimpleNamespace(file_loaders=file_loaders, image_captioning=image_captioning, openai=openai) + vlm = SimpleNamespace( + base_url=vlm_base_url, + model="m", + api_key="k", + timeout=60, + enable_thinking=vlm_enable_thinking, + ) return SimpleNamespace(loader=loader, vlm=vlm) @@ -108,6 +128,35 @@ def test_build_caption_vlm_available_when_endpoint_set_even_if_globally_off() -> assert build_caption_vlm(_config(image_captioning=False, vlm_base_url="http://vlm:8000/v1")) is not None +def test_build_caption_vlm_preserves_enable_thinking() -> None: + vlm = build_caption_vlm( + _config( + image_captioning=False, + vlm_base_url="http://vlm:8000/v1", + vlm_enable_thinking=False, + ) + ) + + assert vlm._enable_thinking is False + + +def test_build_pdf_client_preserves_openai_loader_enable_thinking(monkeypatch: pytest.MonkeyPatch) -> None: + class FakePdfClient: + def __init__(self, vlm, concurrency_limit): + self.vlm = vlm + self.concurrency_limit = concurrency_limit + + import services.inference.parsers.dotsocr as dotsocr + import services.workers.parsers.parser_dispatcher as dispatcher + + monkeypatch.setattr(dotsocr, "DotsOCRPdfClient", FakePdfClient) + monkeypatch.setattr(dispatcher, "_create", lambda _module, _name, **kwargs: kwargs["client"]) + + parser = ParserDispatcher(_config(pdf="DotsOCRLoader", openai_loader_enable_thinking=False))._build_pdf_client() + + assert parser.vlm._enable_thinking is False + + def test_build_eml_wires_nested_email_parser_with_depth_limit(monkeypatch: pytest.MonkeyPatch) -> None: disp = ParserDispatcher(_config()) fallback_parser = _FakeParser() diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index a51aacd64..8fd354314 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -155,7 +155,12 @@ def test_build_contextualizer_factory_uses_global_llm_fallback(tmp_path) -> None (tmp_path / "chunk_contextualizer_tmpl.txt").write_text("Context prompt", encoding="utf-8") cfg = SimpleNamespace( models=SimpleNamespace(llm={}), - llm=SimpleNamespace(base_url="http://llm.example/v1", model="mistral", api_key="llm-key"), + llm=SimpleNamespace( + base_url="http://llm.example/v1", + model="mistral", + api_key="llm-key", + enable_thinking=False, + ), chunker=SimpleNamespace(contextualization_timeout=12, max_concurrent_contextualization=3), semaphore=SimpleNamespace(llm_semaphore=7), paths=SimpleNamespace(prompts_dir=str(tmp_path)), @@ -172,6 +177,7 @@ def test_build_contextualizer_factory_uses_global_llm_fallback(tmp_path) -> None assert contextualizer._llm._endpoint == "http://llm.example/v1" assert contextualizer._llm._model == "mistral" assert contextualizer._llm._api_key == "llm-key" + assert contextualizer._llm._enable_thinking is False # _batch_size drives the per-document loop; _llm.chat is gated by the # injected cluster-wide "llmSemaphore". assert contextualizer._semaphore._name == "llmSemaphore"