diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 9ed83c8c..bdba6660 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Literal +from typing import Any, Literal from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db +from shared.models.schemas.llm_config import LLMConfig from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.app_service import run_retrieval_query from shared.services.retrieval.settings import DEFAULT_TOP_K, VALID_CHUNK_TYPES, normalize_chunk_types @@ -129,12 +130,14 @@ class RetrievalQueryResponse(BaseModel): ) -@router.post("/query", response_model=RetrievalQueryResponse) -async def query_retrieval( +async def execute_retrieval_query( payload: RetrievalQueryRequest, - current_user: CurrentUser = Depends(with_current_user), - db: AsyncSession = Depends(get_db), -): + current_user: CurrentUser, + db: AsyncSession, + *, + llm_config: LLMConfig | None = None, +) -> dict[str, Any]: + """Shared retrieval execution used by v1 and v2 route handlers.""" # Resolve chunk_types: explicit field takes precedence over legacy data_type if payload.chunk_types is not None: resolved_chunk_types = normalize_chunk_types(payload.chunk_types) @@ -166,4 +169,19 @@ async def query_retrieval( threshold=payload.threshold, internal_recall_k=payload.internal_recall_k, use_agentic=payload.use_agentic, + llm_config=llm_config, + ) + + +@router.post("/query", response_model=RetrievalQueryResponse) +async def query_retrieval( + payload: RetrievalQueryRequest, + current_user: CurrentUser = Depends(with_current_user), + db: AsyncSession = Depends(get_db), +): + return await execute_retrieval_query( + payload, + current_user, + db, + llm_config=None, ) diff --git a/apps/api/app/api/v2/routes/retrieval.py b/apps/api/app/api/v2/routes/retrieval.py index e6f18871..6870c72f 100644 --- a/apps/api/app/api/v2/routes/retrieval.py +++ b/apps/api/app/api/v2/routes/retrieval.py @@ -1,5 +1,48 @@ """Retrieval API v2 routes.""" -from app.api.v1.routes.retrieval import router +from __future__ import annotations -__all__ = ["router"] +from app.api.dependencies.current_user import with_current_user +from app.api.v1.routes.retrieval import ( + RetrievalQueryRequest, + RetrievalQueryResponse, + execute_retrieval_query, +) +from app.services.rate_limit.data_structures import CurrentUser +from fastapi import APIRouter, Depends +from pydantic import Field +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.database import get_db +from shared.models.schemas.llm_config import LLMConfig + +router = APIRouter(tags=["Retrieval"]) + + +class RetrievalQueryRequestV2(RetrievalQueryRequest): + """v2 retrieval query request with optional BYOK LLM credentials.""" + + llm_config: LLMConfig | None = Field( + None, + description=( + "Optional bring-your-own-key OpenAI-compatible LLM credentials. " + "Flat {api_key, model, base_url} applies to both channels. " + "Use models.{text,vision} for different model ids on the same " + "endpoint, or text/vision objects for different endpoints. " + "When omitted, server defaults are used." + ), + ) + + +@router.post("/query", response_model=RetrievalQueryResponse) +async def query_retrieval( + payload: RetrievalQueryRequestV2, + current_user: CurrentUser = Depends(with_current_user), + db: AsyncSession = Depends(get_db), +): + return await execute_retrieval_query( + payload, + current_user, + db, + llm_config=payload.llm_config, + ) diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index f082782b..960edb83 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -40,7 +40,7 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) @@ -58,7 +58,8 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> {"role": "system", "content": "you are a helpful assistant"}, {"role": "user", "content": prompt}, ] - resp = get_openai_client().chat_completion( + client, _ = get_text_client() + resp = client.chat_completion( messages=messages, timeout=60, max_tokens=max_tokens, diff --git a/apps/worker/app/services/document_agent/executor/react_loop.py b/apps/worker/app/services/document_agent/executor/react_loop.py index d391c764..a890f702 100644 --- a/apps/worker/app/services/document_agent/executor/react_loop.py +++ b/apps/worker/app/services/document_agent/executor/react_loop.py @@ -235,9 +235,9 @@ def _next_decision(self, round_index: int) -> tuple[ReflexionDecision, ToolResul ) start = time.monotonic() try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client - client = get_openai_client(model=model) + client, model = get_text_client(requested_model=model) raw, usage = client.chat_completion_with_usage( messages=[{"role": "user", "content": prompt}], model=model, diff --git a/apps/worker/app/services/document_agent/planner/planner.py b/apps/worker/app/services/document_agent/planner/planner.py index b7df8671..533cc473 100644 --- a/apps/worker/app/services/document_agent/planner/planner.py +++ b/apps/worker/app/services/document_agent/planner/planner.py @@ -286,9 +286,9 @@ def propose(self) -> tuple[DocumentProfile, ReflexionDecision, ToolResult]: logger.warning("[document_agent] planner png attach failed: {}", exc) try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client - client = get_openai_client(model=model) + client, model = get_vision_client(requested_model=model) raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, diff --git a/apps/worker/app/services/document_agent/structure/page_locate_agent.py b/apps/worker/app/services/document_agent/structure/page_locate_agent.py index ad033f54..7a68f362 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_agent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -328,9 +328,9 @@ def verify_section_page_choice( start = time.monotonic() try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client - client = get_openai_client(model=model) + client, model = get_vision_client(requested_model=model) raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, diff --git a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py index a4c9f2dc..c3826cb2 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_subagent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_subagent.py @@ -299,9 +299,9 @@ def decide_next_action( logger.warning("[page_locate.subagent] planner budget exhausted for title={!r}", title) return None try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client - client = get_openai_client(model=model) + client, model = get_text_client(requested_model=model) raw, usage = client.chat_completion_with_usage( messages=[{"role": "user", "content": prompt}], model=model, diff --git a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py index 33293210..af237901 100644 --- a/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py +++ b/apps/worker/app/services/document_agent/tools/extract_toc_with_boundaries.py @@ -70,7 +70,7 @@ def _vlm_confirm_anchors( budget: Any | None = None, ) -> tuple[list[TocAnchorPage], bool, list[TocEvidence]]: """Phase 1: send all anchor PNGs to VLM, ask which are real TOC starts.""" - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client if not anchor_pages: return [], False, [] @@ -125,7 +125,8 @@ def _vlm_confirm_anchors( return [], True, [] try: - client = get_openai_client(model=model) + client, resolved_model = get_vision_client(requested_model=model) + model = resolved_model or model raw, usage = client.chat_completion_with_usage( messages=messages, model=model, diff --git a/apps/worker/app/services/document_agent/tools/inspect_pages.py b/apps/worker/app/services/document_agent/tools/inspect_pages.py index 8881a9b7..2254dbe4 100644 --- a/apps/worker/app/services/document_agent/tools/inspect_pages.py +++ b/apps/worker/app/services/document_agent/tools/inspect_pages.py @@ -78,9 +78,9 @@ def inspect_pages(ctx: ToolContext, args: dict[str, Any]) -> ToolResult: {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}} ) try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client - client = get_openai_client(model=model) + client, model = get_vision_client(requested_model=model) raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 1ca3fb80..714ba999 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -738,9 +738,9 @@ def propose_shard_plan(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult: if model and ctx.budget.try_reserve("plan", prompt_tokens_est): try: llm_attempted = True - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client - client = get_openai_client(model=model) + client, model = get_text_client(requested_model=model) raw_response, usage = client.chat_completion_with_usage( messages=[{"role": "user", "content": prompt}], model=model, diff --git a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py index bd5c7c04..9c6c87c9 100644 --- a/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py +++ b/apps/worker/app/services/document_agent/tools/vlm_toc_extractor.py @@ -162,7 +162,7 @@ def vlm_extract_toc_batch( BatchTocResult with per-page classification and extracted entries. """ from loguru import logger - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client if not page_pngs: return BatchTocResult( @@ -190,7 +190,8 @@ def vlm_extract_toc_batch( ) start = time.monotonic() - client = get_openai_client(model=model) + client, resolved_model = get_vision_client(requested_model=model) + model = resolved_model or model raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), model=model, diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 04975bbf..4eacc259 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -34,6 +34,7 @@ ) from shared.core.exceptions.domain_exceptions import ValidationException from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.ai.llm_overrides import cleanup_llm_overrides, init_llm_overrides from shared.services.ai.token_tracking import cleanup_token_tracker, init_token_tracker from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service from shared.services.redis.distributed_lock import RedisJobLock @@ -87,6 +88,7 @@ def _run_parse_job( lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") token_usage_dict = init_token_tracker() stage_timing_dict = init_stage_tracker() + init_llm_overrides(JobMetadataHelper.get_llm_config(job_context.job_metadata)) try: prepared_source = prepare_source_file( @@ -180,6 +182,7 @@ def _run_parse_job( "timing_ms": dict(stage_timing_dict), "token_usage": dict(token_usage_dict), } + cleanup_llm_overrides() cleanup_token_tracker() cleanup_stage_tracker() diff --git a/apps/worker/app/services/document_parser/formats/fragment/parser.py b/apps/worker/app/services/document_parser/formats/fragment/parser.py index 1409724b..059604fb 100644 --- a/apps/worker/app/services/document_parser/formats/fragment/parser.py +++ b/apps/worker/app/services/document_parser/formats/fragment/parser.py @@ -16,7 +16,7 @@ from openai.types.chat import ChatCompletionMessageParam from app.services.common.file_utils import path_handle -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.llm_overrides import get_text_client def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str]: @@ -39,7 +39,8 @@ def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str] messages: list[ChatCompletionMessageParam] = [ {"role": "user", "content": title_prompt} ] - generated_title = get_openai_client().chat_completion( + client, _ = get_text_client() + generated_title = client.chat_completion( messages=messages, max_tokens=max_tokens, timeout=30, diff --git a/apps/worker/app/services/document_parser/formats/image/parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py index 726e1523..dd7fdd09 100755 --- a/apps/worker/app/services/document_parser/formats/image/parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -30,10 +30,8 @@ from app.services.common.file_loading import is_remote, load_file_bytes from app.services.common.file_utils import path_handle from shared.services.ai.summary.engine import summarize, transcribe -from shared.services.ai.openai_compatible_client_sync import ( - OpenAICompatibleClientSync, - get_openai_client, -) +from shared.services.ai.llm_overrides import get_vision_client +from shared.services.ai.openai_compatible_client_sync import OpenAICompatibleClientSync MD_IMAGE_PATTERN = r"!\[[^\]]*?\]\((.*?\.(?:png|jpe?g|gif))\)" g_img_lock = threading.Lock() @@ -61,7 +59,8 @@ def perceptual_hash(data: bytes) -> str: def _get_vision_client() -> OpenAICompatibleClientSync: """Create OpenAI-compatible client for vision models, auto-routing by IMAGE_MODEL name.""" image_model = settings.IMAGE_MODEL or "qwen3.6-flash" - return get_openai_client(model=image_model) + client, _ = get_vision_client(requested_model=image_model) + return client def image_bytes_to_base64(img_data: bytes, ext: str) -> str: @@ -151,6 +150,7 @@ def ask_image( image_model = settings.IMAGE_MODEL_MAX or "qwen3.6-flash" if len(urls_) > 0: + client, image_model = get_vision_client(requested_model=image_model) prompt, temperature, top_p, max_tokens = build_prompt( task=task, texts=title_text, query=query, paras={"max_tokens": max_tokens} ) diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 86a20acd..0af79028 100644 --- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -179,7 +179,7 @@ def run_merge_pre_pass( # ── 3. Call LLM directly (bypass df2md — texts is already formatted) ── from shared.services.ai.prompt_service import build_prompt - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client from shared.services.ai.response_process_service import eval_response try: @@ -194,7 +194,8 @@ def run_merge_pre_pass( {"role": "user", "content": prompt}, ] with stage_timer("heading.merge_pre_pass_llm", group_count=len(groups), model_name=model_name): - answer = get_openai_client(model=model_name).chat_completion( + client, model_name = get_text_client(requested_model=model_name) + answer = client.chat_completion( messages=messages, model=model_name, max_tokens=max_tokens, diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py index f99d5329..ea3b8bc8 100755 --- a/apps/worker/app/services/document_parser/structure/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -37,7 +37,7 @@ from shared.services.ai.response_process_service import eval_response # ARQ dependency is removed, use Celery instead -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.llm_overrides import get_text_client # ==================== Helper Functions ==================== @@ -230,7 +230,8 @@ def _is_candidate_id(val): candidate_count=n_candidates, max_tokens=max_tokens, ): - answer = get_openai_client(model=model_name).chat_completion( + client, model_name = get_text_client(requested_model=model_name) + answer = client.chat_completion( messages=messages, model=model_name, max_tokens=max_tokens, diff --git a/apps/worker/app/services/document_parser/structure/toc_parser.py b/apps/worker/app/services/document_parser/structure/toc_parser.py index e3a44cca..88443742 100644 --- a/apps/worker/app/services/document_parser/structure/toc_parser.py +++ b/apps/worker/app/services/document_parser/structure/toc_parser.py @@ -21,7 +21,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.llm_overrides import get_text_client # ==================== Markdown TOC Detection Functions ==================== @@ -230,7 +230,9 @@ def llm_judge_toc_range( model_name=model_name, total_candidates=total_candidates, ): - answer = get_openai_client(model=model_name).chat_completion( + client, resolved_model = get_text_client(requested_model=model_name) + model_name = resolved_model or model_name + answer = client.chat_completion( messages=messages, model=model_name, max_tokens=max_tokens, diff --git a/apps/worker/app/services/document_parser/tables/table_frame_parser.py b/apps/worker/app/services/document_parser/tables/table_frame_parser.py index aa9c4167..38aa92b6 100644 --- a/apps/worker/app/services/document_parser/tables/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_frame_parser.py @@ -15,7 +15,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.llm_overrides import get_text_client from shared.utils.text_utils import remove_duplicates_orderkept @@ -68,7 +68,8 @@ def parse_headers_nonsmart(candidate_frame: pd.DataFrame) -> list[int]: ttl=7200, ) - header_response = get_openai_client().chat_completion( + client, _ = get_text_client() + header_response = client.chat_completion( messages=messages, timeout=60, usage_task="parser.table_detect_headers", diff --git a/apps/worker/app/services/page_memory/fine_hierarchy.py b/apps/worker/app/services/page_memory/fine_hierarchy.py index 361aa072..0e446682 100644 --- a/apps/worker/app/services/page_memory/fine_hierarchy.py +++ b/apps/worker/app/services/page_memory/fine_hierarchy.py @@ -19,7 +19,7 @@ from app.services.page_memory.page_tagger import PageTagResult from app.services.page_memory.skeleton_extractor import SectionSkeleton from app.services.page_memory._utils import page_scope_info -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.ai.llm_overrides import get_text_client from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response @@ -226,7 +226,8 @@ def _run_hierarchy_on_candidates( os.environ.get("NORMOL_MODEL"), ) ) - answer = get_openai_client(model=resolved_model).chat_completion( + client, resolved_model = get_text_client(requested_model=resolved_model) + answer = client.chat_completion( messages=[ {"role": "system", "content": "you are a document structure expert"}, {"role": "user", "content": prompt}, diff --git a/apps/worker/app/services/page_memory/page_assets.py b/apps/worker/app/services/page_memory/page_assets.py index 232bfc83..7e067a71 100644 --- a/apps/worker/app/services/page_memory/page_assets.py +++ b/apps/worker/app/services/page_memory/page_assets.py @@ -109,9 +109,9 @@ def detect_page_assets( ) try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client - client = get_openai_client(model=model_name) + client, model_name = get_vision_client(requested_model=model_name) raw_response, usage = client.chat_completion_with_usage( messages=cast( Any, @@ -829,9 +829,9 @@ def _llm_judge_table_continuity( ) try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client - client = get_openai_client(model=model_name) + client, model_name = get_text_client(requested_model=model_name) raw_response, usage = client.chat_completion_with_usage( messages=[{"role": "user", "content": prompt}], model=model_name, diff --git a/apps/worker/app/services/page_memory/page_tagger.py b/apps/worker/app/services/page_memory/page_tagger.py index 4fe72ac2..4476e6a6 100644 --- a/apps/worker/app/services/page_memory/page_tagger.py +++ b/apps/worker/app/services/page_memory/page_tagger.py @@ -173,9 +173,9 @@ def _tag_text_only( ) try: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_text_client - client = get_openai_client(model=model) + client, model = get_text_client(requested_model=model) raw_response, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": prompt}]), model=model, @@ -388,9 +388,10 @@ def _tag_vlm_titles( }, ] - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import get_vision_client - client = get_openai_client(model=model) + client, resolved_model = get_vision_client(requested_model=model) + model = resolved_model or model for attempt in range(_MAX_JSON_RETRIES + 1): try: diff --git a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py index 7725d396..78f13946 100644 --- a/apps/worker/tests/contract/test_page_memory_asset_java_contract.py +++ b/apps/worker/tests/contract/test_page_memory_asset_java_contract.py @@ -46,7 +46,7 @@ def chat_completion_with_usage(self, **kwargs): import shared.services.ai.openai_compatible_client_sync as client_mod monkeypatch.setattr( - client_mod, "get_openai_client", lambda model=None: _FakeClient() + client_mod, "get_openai_client", lambda model=None, **_kwargs: _FakeClient() ) image_path = tmp_path / "page.png" diff --git a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py index c46eac98..4eac1d24 100644 --- a/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py +++ b/apps/worker/tests/contract/test_page_memory_fine_hierarchy_contract.py @@ -98,8 +98,11 @@ def test_refine_fat_leaf_skeletons_excludes_next_section_start_when_unordered( monkeypatch.setattr( fine_hierarchy, - "get_openai_client", - lambda model=None: _FakeClient([{"id": 1, "level": 1}]), + "get_text_client", + lambda requested_model=None: ( + _FakeClient([{"id": 1, "level": 1}]), + requested_model, + ), ) refined = fine_hierarchy.refine_fat_leaf_skeletons( @@ -152,14 +155,17 @@ def test_refine_fat_leaf_skeletons_uses_page_memory_prompt_without_demoting_sibl monkeypatch.setattr( fine_hierarchy, - "get_openai_client", - lambda model=None: _FakeClient( - [ - {"id": 1, "level": 1}, - {"id": 2, "level": 1}, - {"id": 3, "level": 2}, - {"id": 4, "level": 1}, - ] + "get_text_client", + lambda requested_model=None: ( + _FakeClient( + [ + {"id": 1, "level": 1}, + {"id": 2, "level": 1}, + {"id": 3, "level": 2}, + {"id": 4, "level": 1}, + ] + ), + requested_model, ), ) diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index 21304323..a5e2edb2 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -410,7 +410,7 @@ def chat_completion_with_usage(self, **kwargs): import shared.services.ai.openai_compatible_client_sync as client_mod monkeypatch.setattr( - client_mod, "get_openai_client", lambda model=None: _FakeClient() + client_mod, "get_openai_client", lambda model=None, **_kwargs: _FakeClient() ) img = tmp_path / "page-231.png" diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index 4326ba43..2a948f79 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -5,6 +5,8 @@ from pydantic import BaseModel, ConfigDict, Field +from shared.models.schemas.llm_config import LLMConfig + class WebhookConfig(BaseModel): """Webhook configuration.""" @@ -85,6 +87,17 @@ class JobCreate(JobCreateBase): class JobCreateV2(JobCreateBase): """Public v2 request payload for creating a job.""" + llm_config: Optional[LLMConfig] = Field( + None, + description=( + "Optional bring-your-own-key OpenAI-compatible LLM credentials. " + "Flat {api_key, model, base_url} applies to both channels. " + "Use models.{text,vision} for different model ids on the same " + "endpoint, or text/vision objects for different endpoints. " + "When omitted, server defaults are used." + ), + ) + class JobResponse(BaseModel): """Job-creation response.""" diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index 30934dc5..4efb808b 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -4,8 +4,10 @@ from pydantic import BaseModel, ConfigDict, Field +from shared.models.schemas.llm_config import LLMConfig, parse_llm_config from shared.models.schemas.page_memory_config import PageMemoryConfig from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace +from shared.utils.security_utils import mask_api_key class JobMetadataBase(BaseModel): @@ -18,6 +20,9 @@ class JobMetadataBase(BaseModel): parsing_params: Optional[Dict[str, Any]] = Field( None, description="Parsing parameters" ) + llm_config: Optional[Dict[str, Any]] = Field( + None, description="BYOK OpenAI-compatible LLM credentials (v2)" + ) data_id: Optional[str] = Field(None, description="User-defined ID") webhook: Optional[Dict[str, Any]] = Field(None, description="Webhook configuration") document_metadata: Optional[Dict[str, Any]] = Field( @@ -67,6 +72,13 @@ def create_from_request( resolved_page_memory_config = page_memory_config.to_dict() else: resolved_page_memory_config = page_memory_config + # v2-only field; getattr keeps v1 JobCreate (no llm_config) safe. + raw_llm_config = getattr(request, "llm_config", None) + llm_config_payload: Dict[str, Any] | None = None + if isinstance(raw_llm_config, LLMConfig): + llm_config_payload = raw_llm_config.model_dump() + elif isinstance(raw_llm_config, dict): + llm_config_payload = dict(raw_llm_config) metadata = { "original_request": _dump_public_request(request), "api_version": api_version, @@ -81,6 +93,8 @@ def create_from_request( "data_id": request.data_id, "webhook": request.webhook.model_dump() if request.webhook else None, } + if llm_config_payload is not None: + metadata["llm_config"] = llm_config_payload if resolved_page_memory_config is not None: metadata["page_memory_config"] = resolved_page_memory_config metadata.update(kwargs) @@ -239,8 +253,40 @@ def get_webhook(metadata: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Return the webhook configuration from metadata.""" return JobMetadataHelper.get_field(metadata, "webhook") + @staticmethod + def get_llm_config(metadata: Optional[Dict[str, Any]]) -> LLMConfig | None: + """Return the BYOK LLM config stored in metadata, if any.""" + raw = JobMetadataHelper.get_field(metadata, "llm_config", None) + return parse_llm_config(raw) + + +def _mask_llm_config_in_request(payload: Dict[str, Any]) -> Dict[str, Any]: + """Redact api_key values inside llm_config for public request snapshots.""" + llm_config = payload.get("llm_config") + if not isinstance(llm_config, dict): + return payload + + masked = dict(payload) + masked_llm: Dict[str, Any] = dict(llm_config) + if isinstance(masked_llm.get("api_key"), str): + masked_llm["api_key"] = mask_api_key(masked_llm["api_key"]) + for slot in ("text", "vision"): + provider = masked_llm.get(slot) + if isinstance(provider, dict): + provider_copy = dict(provider) + if "api_key" in provider_copy: + provider_copy["api_key"] = mask_api_key( + provider_copy.get("api_key") + if isinstance(provider_copy.get("api_key"), str) + else None + ) + masked_llm[slot] = provider_copy + masked["llm_config"] = masked_llm + return masked + def _dump_public_request(request) -> Dict[str, Any]: """Dump declared public request fields without hidden compatibility extras.""" extra_fields = getattr(request, "model_extra", None) or {} - return request.model_dump(exclude=set(extra_fields)) + payload = request.model_dump(exclude=set(extra_fields)) + return _mask_llm_config_in_request(payload) diff --git a/packages/shared-python/shared/models/schemas/llm_config.py b/packages/shared-python/shared/models/schemas/llm_config.py new file mode 100644 index 00000000..cdcc3d4b --- /dev/null +++ b/packages/shared-python/shared/models/schemas/llm_config.py @@ -0,0 +1,174 @@ +"""Bring-your-own-key (BYOK) OpenAI-compatible LLM credentials.""" + +from __future__ import annotations + +from typing import Any, Optional + +from pydantic import BaseModel, Field, model_validator + + +class LLMProviderConfig(BaseModel): + """Credentials for one OpenAI-compatible provider endpoint.""" + + api_key: str = Field(..., min_length=1, description="Provider API key") + model: str = Field(..., min_length=1, description="Model identifier") + base_url: str = Field( + ..., + min_length=1, + description="OpenAI-compatible base URL (e.g. https://api.openai.com/v1)", + ) + + +class LLMModelsConfig(BaseModel): + """Per-channel model ids that share root api_key / base_url.""" + + text: Optional[str] = Field(None, min_length=1, description="Text / planning model id") + vision: Optional[str] = Field(None, min_length=1, description="Vision / VLM model id") + + +class LLMConfig(BaseModel): + """OpenAI-compatible BYOK credentials (flat root + optional channel overrides). + + Happy path (one multimodal model for both channels):: + + {"api_key": "...", "model": "gpt-4o", "base_url": "https://api.openai.com/v1"} + + Same endpoint, different models per channel:: + + { + "api_key": "...", + "base_url": "https://api.openai.com/v1", + "models": {"text": "gpt-4o-mini", "vision": "gpt-4o"} + } + + Different endpoints per channel:: + + { + "text": {"api_key": "...", "model": "...", "base_url": "..."}, + "vision": {"api_key": "...", "model": "...", "base_url": "..."} + } + + Semantics: + - root ``api_key`` + ``base_url`` with ``model`` and/or ``models`` -> shared auth + - ``models.`` wins over root ``model`` for that channel + - ``text`` / ``vision`` objects fully replace the root for that channel + - a channel with no resolved model keeps server defaults + """ + + api_key: Optional[str] = Field(None, min_length=1, description="Default provider API key") + model: Optional[str] = Field( + None, + min_length=1, + description="Default model for both channels (overridden by models.*)", + ) + base_url: Optional[str] = Field( + None, + min_length=1, + description="Default OpenAI-compatible base URL", + ) + models: Optional[LLMModelsConfig] = Field( + None, + description="Per-channel model ids sharing root api_key / base_url", + ) + text: Optional[LLMProviderConfig] = Field( + None, + description="Text / planning credentials (replaces root for text channel)", + ) + vision: Optional[LLMProviderConfig] = Field( + None, + description="Vision / VLM credentials (replaces root for vision channel)", + ) + + @model_validator(mode="after") + def _validate_shape(self) -> "LLMConfig": + has_api_key = self.api_key is not None + has_base_url = self.base_url is not None + if has_api_key != has_base_url: + raise ValueError("llm_config api_key and base_url must be set together") + + has_auth = has_api_key and has_base_url + has_models = self.models is not None and ( + self.models.text is not None or self.models.vision is not None + ) + if self.models is not None and not has_models: + raise ValueError("llm_config.models requires at least one of text or vision") + if self.model is not None and not has_auth: + raise ValueError("llm_config.model requires api_key and base_url") + if has_models and not has_auth: + raise ValueError("llm_config.models requires api_key and base_url") + if has_auth and self.model is None and not has_models: + raise ValueError( + "llm_config with api_key/base_url requires model and/or models" + ) + + if ( + not has_auth + and self.text is None + and self.vision is None + ): + raise ValueError( + "llm_config requires root credentials and/or text/vision overrides" + ) + return self + + def _channel_model(self, channel: str) -> str | None: + if self.models is not None: + named = getattr(self.models, channel) + if isinstance(named, str) and named: + return named + return self.model + + def _root_provider_for(self, channel: str) -> LLMProviderConfig | None: + if self.api_key is None or self.base_url is None: + return None + model = self._channel_model(channel) + if model is None: + return None + return LLMProviderConfig( + api_key=self.api_key, + model=model, + base_url=self.base_url, + ) + + def text_effective(self) -> LLMProviderConfig | None: + """Return the text-channel config, or None to keep server defaults.""" + return self.text if self.text is not None else self._root_provider_for("text") + + def vision_effective(self) -> LLMProviderConfig | None: + """Return the vision-channel config, or None to keep server defaults.""" + return ( + self.vision if self.vision is not None else self._root_provider_for("vision") + ) + + def masked_dump(self) -> dict[str, Any]: + """Serialize with api_key values redacted for snapshots / responses.""" + from shared.utils.security_utils import mask_api_key + + def _mask_provider(provider: LLMProviderConfig | None) -> dict[str, Any] | None: + if provider is None: + return None + return { + "api_key": mask_api_key(provider.api_key), + "model": provider.model, + "base_url": provider.base_url, + } + + return { + "api_key": mask_api_key(self.api_key) if self.api_key else None, + "model": self.model, + "base_url": self.base_url, + "models": self.models.model_dump() if self.models is not None else None, + "text": _mask_provider(self.text), + "vision": _mask_provider(self.vision), + } + + +def parse_llm_config(value: Any) -> LLMConfig | None: + """Parse a raw mapping / LLMConfig into a validated LLMConfig, or None.""" + if value is None: + return None + if isinstance(value, LLMConfig): + return value + if isinstance(value, dict): + return LLMConfig.model_validate(value) + return None diff --git a/packages/shared-python/shared/services/ai/llm_overrides.py b/packages/shared-python/shared/services/ai/llm_overrides.py new file mode 100644 index 00000000..bc5671ec --- /dev/null +++ b/packages/shared-python/shared/services/ai/llm_overrides.py @@ -0,0 +1,160 @@ +"""Request-scoped BYOK LLM credential overrides. + +In the gevent worker, child greenlets (GeventPool.spawn) do NOT inherit +``ContextVar`` or ``threading.local`` from the parent. We therefore use +a module-level dict keyed by the *root* greenlet id of the current parse +task — the same pattern as ``token_tracking``. + +For the asyncio retrieval path a ``ContextVar`` is used instead, which +propagates naturally across ``await`` and ``asyncio.to_thread``. + +``get_current_llm_overrides()`` checks the greenlet dict first, then the +ContextVar, so both runtimes share one lookup API. +""" + +from __future__ import annotations + +import threading +from contextvars import ContextVar, Token +from typing import Any, Optional + +from shared.models.schemas.llm_config import LLMConfig, LLMProviderConfig, parse_llm_config + +# Greenlet-root keyed store (worker / gevent). +_overrides: dict[int, LLMConfig] = {} +_root_ids: dict[int, int] = {} +_lock = threading.Lock() + +# Async ContextVar store (API retrieval). +_async_overrides: ContextVar[LLMConfig | None] = ContextVar( + "llm_overrides_async", + default=None, +) + +ResolvedCredentials = tuple[str | None, str | None, str | None] +# (model, api_key, api_url) + + +def _current_greenlet_id() -> int: + try: + import gevent + + return id(gevent.getcurrent()) + except ImportError: + return threading.get_ident() + + +def _find_root_id() -> int | None: + """Walk up the greenlet parent chain to find a registered root id.""" + gid = _current_greenlet_id() + if gid in _overrides: + return gid + if gid in _root_ids: + return _root_ids[gid] + try: + import gevent + except ImportError: + # gevent is worker-only; without it there is no parent chain to walk. + return None + + g = gevent.getcurrent() + while g is not None: + pid = id(g) + if pid in _overrides: + _root_ids[gid] = pid + return pid + g = getattr(g, "parent", None) + return None + + +def init_llm_overrides(config: LLMConfig | dict[str, Any] | None) -> LLMConfig | None: + """Register BYOK overrides for the current parse-task root greenlet. + + Must be called from the root greenlet of the task (same scope as + ``init_token_tracker``). Returns the parsed config, or None when + no override is active. + """ + parsed = parse_llm_config(config) + if parsed is None: + return None + gid = _current_greenlet_id() + with _lock: + _overrides[gid] = parsed + return parsed + + +def cleanup_llm_overrides() -> None: + """Remove the override for the current greenlet. Call after parsing.""" + gid = _current_greenlet_id() + with _lock: + _overrides.pop(gid, None) + stale = [k for k, v in _root_ids.items() if v == gid] + for k in stale: + del _root_ids[k] + + +def set_llm_overrides_async( + config: LLMConfig | dict[str, Any] | None, +) -> Token[LLMConfig | None]: + """Set BYOK overrides for the current asyncio task. Returns a reset token.""" + parsed = parse_llm_config(config) + return _async_overrides.set(parsed) + + +def reset_llm_overrides_async(token: Token[LLMConfig | None]) -> None: + """Reset the async override ContextVar using the token from set_*.""" + _async_overrides.reset(token) + + +def get_current_llm_overrides() -> LLMConfig | None: + """Return the active BYOK config for this task, if any. + + Checks the gevent greenlet-root store first, then the async ContextVar. + """ + root = _find_root_id() + if root is not None: + return _overrides.get(root) + return _async_overrides.get() + + +def _resolve_provider( + provider: LLMProviderConfig | None, + requested_model: str | None, +) -> ResolvedCredentials: + if provider is None: + return (requested_model, None, None) + return (provider.model, provider.api_key, provider.base_url) + + +def resolve_text(requested_model: str | None = None) -> ResolvedCredentials: + """Resolve (model, api_key, api_url) for a text LLM call.""" + overrides = get_current_llm_overrides() + if overrides is None: + return (requested_model, None, None) + return _resolve_provider(overrides.text_effective(), requested_model) + + +def resolve_vision(requested_model: str | None = None) -> ResolvedCredentials: + """Resolve (model, api_key, api_url) for a vision/VLM call.""" + overrides = get_current_llm_overrides() + if overrides is None: + return (requested_model, None, None) + return _resolve_provider(overrides.vision_effective(), requested_model) + + +def get_text_client(requested_model: Optional[str] = None): + """Return ``(client, effective_model)`` for a text LLM call.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + model, api_key, api_url = resolve_text(requested_model) + client = get_openai_client(model=model, api_key=api_key, api_url=api_url) + return client, model + + +def get_vision_client(requested_model: Optional[str] = None): + """Return ``(client, effective_model)`` for a vision/VLM call.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + + model, api_key, api_url = resolve_vision(requested_model) + client = get_openai_client(model=model, api_key=api_key, api_url=api_url) + return client, model diff --git a/packages/shared-python/shared/services/ai/summary/engine.py b/packages/shared-python/shared/services/ai/summary/engine.py index eaedd05c..141520f6 100644 --- a/packages/shared-python/shared/services/ai/summary/engine.py +++ b/packages/shared-python/shared/services/ai/summary/engine.py @@ -108,6 +108,7 @@ def _call_llm( budget: Any | None, budget_pool: str, budget_stage: str | None, + channel: Literal["text", "vision"] = "text", ) -> Any | None: """One text-or-vision call with budget accounting and a single JSON retry. @@ -115,7 +116,11 @@ def _call_llm( failure / exhausted budget. Budget is reserved before the call, committed on success, refunded on failure — matching the prior per-caller bookkeeping but in one place. + + ``channel`` selects BYOK text vs vision credentials when overrides are active. """ + from shared.services.ai.llm_overrides import resolve_text, resolve_vision + content_parts: list[dict[str, Any]] = [{"type": "text", "text": prompt}] for path in image_paths: img_b64 = _read_image_b64(path) @@ -142,12 +147,23 @@ def _call_llm( if expect_json: api_kwargs["response_format"] = {"type": "json_object"} - client = _client_mod.get_openai_client(model=model) + resolve = resolve_vision if channel == "vision" else resolve_text + effective_model, api_key, api_url = resolve(model) + if not effective_model: + if budget is not None: + budget.refund(budget_pool, est=est, stage=budget_stage) + return None + + client = _client_mod.get_openai_client( + model=effective_model, + api_key=api_key, + api_url=api_url, + ) for attempt in range(_MAX_JSON_RETRIES + 1): try: raw, usage = client.chat_completion_with_usage( messages=cast(Any, [{"role": "user", "content": content_parts}]), - model=model, + model=effective_model, temperature=temperature, top_p=top_p, max_tokens=max_tokens, @@ -337,6 +353,7 @@ def _summarize_body( budget=budget, budget_pool=budget_pool, budget_stage=budget_stage, + channel="vision", ) else: # Text summary: the shared summary-full prompt with deterministic lang lock. @@ -366,6 +383,7 @@ def _summarize_body( budget=budget, budget_pool="plan", budget_stage=None, + channel="text", ) if not isinstance(parsed, dict): @@ -416,6 +434,7 @@ def _summarize_asset( budget=budget, budget_pool="plan", budget_stage=None, + channel="text", ) if isinstance(raw, str) and raw.strip(): return _parse_linesplit_asset(raw, asset_title_hint) @@ -444,6 +463,7 @@ def _summarize_asset( budget=budget, budget_pool=budget_pool, budget_stage=budget_stage, + channel="vision", ) if isinstance(parsed, dict): return AssetSummary( @@ -492,6 +512,7 @@ def transcribe( budget=budget, budget_pool=budget_pool, budget_stage=budget_stage, + channel="vision", ) if isinstance(parsed, dict): return str(parsed.get("text", "")).strip() diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 1e061510..83ce65b7 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -4,6 +4,7 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.models.schemas.llm_config import LLMConfig from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.execution.plan import ( run_retrieval_query as execute_retrieval_query, @@ -31,6 +32,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + llm_config: LLMConfig | None = None, ) -> dict[str, Any]: return await execute_retrieval_query( db=db, @@ -49,4 +51,5 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + llm_config=llm_config, ) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index 99e012e0..ed9f9f59 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -6,6 +6,11 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from shared.models.schemas.llm_config import LLMConfig +from shared.services.ai.llm_overrides import ( + reset_llm_overrides_async, + set_llm_overrides_async, +) from shared.services.retrieval.cache_service import ( get_cached_retrieval_query_result, set_cached_retrieval_query_result, @@ -38,6 +43,7 @@ async def run_retrieval_query( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + llm_config: LLMConfig | None = None, ) -> dict[str, Any]: """Run retrieval through the plan module.""" return await RetrievalExecutionPlan( @@ -58,6 +64,7 @@ async def run_retrieval_query( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + llm_config=llm_config, ) ).execute() @@ -68,7 +75,14 @@ def __init__(self, request: RetrievalQuery) -> None: async def execute(self) -> dict[str, Any]: request = self.request + override_token = set_llm_overrides_async(request.llm_config) + + try: + return await self._execute_with_overrides(request) + finally: + reset_llm_overrides_async(override_token) + async def _execute_with_overrides(self, request: RetrievalQuery) -> dict[str, Any]: # TODO(intent-step): Insert Intent Understanding step here. # Before any retrieval runs, parse `request.query` with LLM + # KG overview + section tree to extract structured navigation diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py index a7094b88..c226a9f6 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -5,6 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.models.schemas.llm_config import LLMConfig from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.execution.route_types import RetrievalRouteContext from shared.services.retrieval.settings import ( @@ -30,6 +31,7 @@ class RetrievalQuery: threshold: float = 0.0 internal_recall_k: int | None = None use_agentic: bool | None = None + llm_config: LLMConfig | None = None @classmethod def from_parameters( @@ -51,6 +53,7 @@ def from_parameters( threshold: float = 0.0, internal_recall_k: int | None = None, use_agentic: bool | None = None, + llm_config: LLMConfig | None = None, ) -> "RetrievalQuery": return cls( db=db, @@ -69,9 +72,17 @@ def from_parameters( threshold=threshold, internal_recall_k=internal_recall_k, use_agentic=use_agentic, + llm_config=llm_config, ) def build_cache_extra(self) -> dict[str, Any]: + text_model: str | None = None + vision_model: str | None = None + if self.llm_config is not None: + text_provider = self.llm_config.text_effective() + vision_provider = self.llm_config.vision_effective() + text_model = text_provider.model if text_provider is not None else None + vision_model = vision_provider.model if vision_provider is not None else None return { "chunk_types": sorted(self.chunk_types) if self.chunk_types else None, "signal_paths": self.signal_paths, @@ -83,6 +94,8 @@ def build_cache_extra(self) -> dict[str, Any]: "internal_recall_k": self.internal_recall_k, "use_agentic": self.use_agentic, "decomposition_enabled": True, + "llm_text_model": text_model, + "llm_vision_model": vision_model, } def resolve_allowed_chunk_types(self) -> set[str] | None: diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index 9ee34401..8c8cb627 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -13,6 +13,7 @@ from loguru import logger from shared.core.config import settings +from shared.services.ai.llm_overrides import get_current_llm_overrides # LLMFn accepts either a plain string or a list of ChatCompletionMessageParam LLMFnInput = Union[str, Sequence[dict[str, Any]]] @@ -29,6 +30,8 @@ def _has_llm_credentials() -> bool: """Check whether at least one LLM provider is configured.""" + if get_current_llm_overrides() is not None: + return True if getattr(settings, 'LLM_MOCK_ENABLED', False): return True if getattr(settings, 'DS_KEY', ''): @@ -44,6 +47,11 @@ def _has_llm_credentials() -> bool: def _resolve_default_model() -> str: """Pick a model name that matches the configured LLM provider.""" + overrides = get_current_llm_overrides() + if overrides is not None: + provider = overrides.text_effective() + if provider is not None: + return provider.model if getattr(settings, 'DS_KEY', ''): return 'deepseek-v4-flash' if getattr(settings, 'ALI_API_KEYS', ''): @@ -56,6 +64,11 @@ def _resolve_default_model() -> str: def _resolve_planner_model(*, thinking: bool) -> str: + overrides = get_current_llm_overrides() + if overrides is not None: + provider = overrides.text_effective() + if provider is not None: + return provider.model configured = getattr(settings, 'RETRIEVAL_PLANNER_MODEL', '') or '' if configured: return configured @@ -70,6 +83,29 @@ def _resolve_planner_model(*, thinking: bool) -> str: return getattr(settings, 'NORMOL_MODEL', None) or 'deepseek-v4-flash' +def _resolve_vlm_model(model: str | None = None) -> str: + overrides = get_current_llm_overrides() + if overrides is not None: + provider = overrides.vision_effective() + if provider is not None: + return provider.model + return model or getattr(settings, 'IMAGE_MODEL', '') or 'qwen3.6-flash' + + +def _build_client_for_channel(*, channel: str, model: str): + """Build an OpenAI-compatible client, honoring active BYOK overrides.""" + from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.services.ai.llm_overrides import resolve_text, resolve_vision + + resolve = resolve_vision if channel == 'vision' else resolve_text + effective_model, api_key, api_url = resolve(model) + return get_openai_client( + model=effective_model, + api_key=api_key, + api_url=api_url, + ), effective_model + + def create_retrieval_llm_fn( *, model: str | None = None, @@ -108,9 +144,10 @@ def create_retrieval_llm_fn( ) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=effective_model) + client, resolved_model = _build_client_for_channel( + channel='text', + model=effective_model, + ) current_llm_usage.set(None) kwargs: dict[str, Any] = {} @@ -124,7 +161,7 @@ async def llm_fn(prompt: LLMFnInput) -> str: result, usage = await asyncio.to_thread( client.chat_completion_with_usage, cast(Any, prompt), - model=effective_model, + model=resolved_model, temperature=effective_temperature, max_tokens=effective_max_tokens, **kwargs, @@ -149,14 +186,15 @@ def create_retrieval_planner_fn( effective_model = model or _resolve_planner_model(thinking=thinking) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=effective_model) + client, resolved_model = _build_client_for_channel( + channel='text', + model=effective_model, + ) current_llm_usage.set(None) result, usage = await asyncio.to_thread( client.chat_completion_with_usage, cast(Any, prompt), - model=effective_model, + model=resolved_model, temperature=0.0, max_tokens=max_tokens, ) @@ -181,23 +219,22 @@ def create_retrieval_vlm_fn( ``create_retrieval_llm_fn`` — callers pass either a plain string or a list of ChatCompletionMessageParam (including image_url parts). """ - from shared.core.config import settings - - effective_model = model or getattr(settings, 'IMAGE_MODEL', '') or 'qwen3.6-flash' + effective_model = _resolve_vlm_model(model) if not _has_llm_credentials(): logger.debug('retrieval: no LLM credentials for VLM, image-aware answering disabled') return None async def vlm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client - - client = get_openai_client(model=effective_model) + client, resolved_model = _build_client_for_channel( + channel='vision', + model=effective_model, + ) current_llm_usage.set(None) result, usage = await asyncio.to_thread( client.chat_completion_with_usage, cast(Any, prompt), - model=effective_model, + model=resolved_model, temperature=temperature, max_tokens=max_tokens, ) diff --git a/packages/shared-python/shared/tests/test_llm_config.py b/packages/shared-python/shared/tests/test_llm_config.py new file mode 100644 index 00000000..70d7eb6d --- /dev/null +++ b/packages/shared-python/shared/tests/test_llm_config.py @@ -0,0 +1,150 @@ +"""Unit tests for BYOK LLMConfig resolution.""" + +from __future__ import annotations + +import os + +import pytest +from pydantic import ValidationError + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.models.schemas.llm_config import LLMConfig, LLMModelsConfig, LLMProviderConfig + + +def _creds( + model: str = "gpt-4o", + *, + api_key: str = "sk-test", + base_url: str = "https://api.openai.com/v1", +) -> LLMProviderConfig: + return LLMProviderConfig(api_key=api_key, model=model, base_url=base_url) + + +def test_flat_root_applies_to_both_channels() -> None: + cfg = LLMConfig( + api_key="sk-root", + model="gpt-4o", + base_url="https://api.openai.com/v1", + ) + assert cfg.text_effective() is not None + assert cfg.vision_effective() is not None + assert cfg.text_effective().model == "gpt-4o" + assert cfg.vision_effective().api_key == "sk-root" + + +def test_models_map_same_endpoint_different_models() -> None: + cfg = LLMConfig( + api_key="sk-root", + base_url="https://api.openai.com/v1", + models=LLMModelsConfig(text="gpt-4o-mini", vision="gpt-4o"), + ) + assert cfg.text_effective().model == "gpt-4o-mini" + assert cfg.vision_effective().model == "gpt-4o" + assert cfg.text_effective().api_key == "sk-root" + assert cfg.vision_effective().base_url == "https://api.openai.com/v1" + + +def test_models_map_partial_leaves_other_channel_default() -> None: + cfg = LLMConfig( + api_key="sk-root", + base_url="https://api.openai.com/v1", + models=LLMModelsConfig(text="gpt-4o-mini"), + ) + assert cfg.text_effective().model == "gpt-4o-mini" + assert cfg.vision_effective() is None + + +def test_models_overrides_root_model_per_channel() -> None: + cfg = LLMConfig( + api_key="sk-root", + model="gpt-4o-mini", + base_url="https://api.openai.com/v1", + models=LLMModelsConfig(vision="gpt-4o"), + ) + assert cfg.text_effective().model == "gpt-4o-mini" + assert cfg.vision_effective().model == "gpt-4o" + + +def test_text_only_leaves_vision_on_defaults() -> None: + cfg = LLMConfig(text=_creds("text-model")) + assert cfg.text_effective().model == "text-model" + assert cfg.vision_effective() is None + + +def test_vision_only_leaves_text_on_defaults() -> None: + cfg = LLMConfig(vision=_creds("vlm")) + assert cfg.text_effective() is None + assert cfg.vision_effective().model == "vlm" + + +def test_two_different_endpoints() -> None: + cfg = LLMConfig( + text=_creds("gpt-4o-mini", base_url="https://api.openai.com/v1"), + vision=_creds( + "qwen-vl-max", + api_key="sk-ali", + base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", + ), + ) + assert cfg.text_effective().base_url == "https://api.openai.com/v1" + assert cfg.vision_effective().base_url.endswith("/compatible-mode/v1") + assert cfg.vision_effective().model == "qwen-vl-max" + + +def test_channel_replaces_root() -> None: + cfg = LLMConfig( + api_key="sk-root", + model="shared", + base_url="https://api.openai.com/v1", + text=_creds("text-only"), + vision=_creds("vision-only", base_url="https://other.example/v1"), + ) + assert cfg.text_effective().model == "text-only" + assert cfg.vision_effective().model == "vision-only" + assert cfg.vision_effective().base_url == "https://other.example/v1" + + +def test_root_plus_text_override() -> None: + cfg = LLMConfig( + api_key="sk-root", + model="shared", + base_url="https://api.openai.com/v1", + text=_creds("text-only"), + ) + assert cfg.text_effective().model == "text-only" + assert cfg.vision_effective().model == "shared" + + +def test_partial_root_rejected() -> None: + with pytest.raises(ValidationError, match="api_key and base_url must be set together"): + LLMConfig(api_key="sk-only") + + +def test_auth_without_model_rejected() -> None: + with pytest.raises(ValidationError, match="requires model and/or models"): + LLMConfig(api_key="sk-root", base_url="https://api.openai.com/v1") + + +def test_empty_config_rejected() -> None: + with pytest.raises(ValidationError, match="root credentials and/or text/vision"): + LLMConfig() + + +def test_masked_dump_masks_root_and_channels() -> None: + cfg = LLMConfig( + api_key="sk-root", + model="gpt-4o", + base_url="https://api.openai.com/v1", + vision=_creds("vlm", api_key="sk-vision"), + ) + dump = cfg.masked_dump() + assert dump["api_key"] != "sk-root" + assert dump["model"] == "gpt-4o" + assert dump["vision"]["api_key"] != "sk-vision" + assert dump["text"] is None diff --git a/scripts/smoke_byok.sh b/scripts/smoke_byok.sh new file mode 100755 index 00000000..104848f7 --- /dev/null +++ b/scripts/smoke_byok.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +# Local BYOK smoke checks against a running KNOWHERE API on :5005. +# Usage: +# export KNOWHERE_API_KEY=kh_... +# ./scripts/smoke_byok.sh +set -euo pipefail + +BASE_URL="${KNOWHERE_BASE_URL:-http://localhost:5005}" +API_KEY="${KNOWHERE_API_KEY:?Set KNOWHERE_API_KEY to a local API key}" + +auth_hdr=(-H "Authorization: Bearer ${API_KEY}" -H "Content-Type: application/json") + +echo "== health ==" +curl -fsS "${BASE_URL}/health" | head -c 200; echo + +echo "== OpenAPI: v2 jobs has llm_config, v1 does not ==" +python3 - <<'PY' +import json, os, urllib.request +base = os.environ.get("KNOWHERE_BASE_URL", "http://localhost:5005") +with urllib.request.urlopen(f"{base}/openapi.json") as resp: + schema = json.load(resp) +paths = schema["paths"] +v1 = paths.get("/api/v1/jobs", paths.get("/v1/jobs", {})) +v2 = paths.get("/api/v2/jobs", paths.get("/v2/jobs", {})) +# FastAPI root_path may strip /api from openapi paths — try both +def body_props(op): + if not op: + return set() + post = op.get("post") or {} + body = ((post.get("requestBody") or {}).get("content") or {}).get("application/json") or {} + s = body.get("schema") or {} + if "$ref" in s: + name = s["$ref"].rsplit("/", 1)[-1] + s = schema["components"]["schemas"].get(name, {}) + return set((s.get("properties") or {}).keys()) + +# Prefer component schemas when available +comps = schema.get("components", {}).get("schemas", {}) +v1_keys = set((comps.get("JobCreate") or {}).get("properties", {}).keys()) +v2_keys = set((comps.get("JobCreateV2") or {}).get("properties", {}).keys()) +r1 = set((comps.get("RetrievalQueryRequest") or {}).get("properties", {}).keys()) +r2 = set((comps.get("RetrievalQueryRequestV2") or {}).get("properties", {}).keys()) +print("JobCreate llm_config:", "llm_config" in v1_keys) +print("JobCreateV2 llm_config:", "llm_config" in v2_keys) +print("RetrievalQueryRequest llm_config:", "llm_config" in r1) +print("RetrievalQueryRequestV2 llm_config:", "llm_config" in r2) +assert "llm_config" not in v1_keys +assert "llm_config" in v2_keys +assert "llm_config" not in r1 +assert "llm_config" in r2 +print("OpenAPI BYOK surface OK") +PY + +echo "== v2 jobs: reject empty llm_config object ==" +code=$(curl -sS -o /tmp/byok_empty.json -w '%{http_code}' "${auth_hdr[@]}" \ + -d '{"source_type":"url","source_url":"https://example.com/a.pdf","llm_config":{}}' \ + "${BASE_URL}/api/v2/jobs") +echo "HTTP $code" +cat /tmp/byok_empty.json | head -c 400; echo +test "$code" = "422" || test "$code" = "400" + +echo "== v2 jobs: accept flat llm_config (multimodal shorthand) shape ==" +# Expect waiting-file or pending-ish success, not 422 +code=$(curl -sS -o /tmp/byok_ok.json -w '%{http_code}' "${auth_hdr[@]}" \ + -d '{ + "source_type":"url", + "source_url":"https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", + "file_name":"dummy.pdf", + "data_id":"byok-smoke-flat", + "llm_config":{ + "api_key":"sk-smoke-test-key-not-real", + "model":"gpt-4o", + "base_url":"https://api.openai.com/v1" + } + }' \ + "${BASE_URL}/api/v2/jobs") +echo "HTTP $code" +cat /tmp/byok_ok.json | head -c 600; echo +test "$code" = "200" || test "$code" = "201" + +JOB_ID=$(python3 -c 'import json; print(json.load(open("/tmp/byok_ok.json")).get("job_id",""))') +echo "job_id=${JOB_ID}" + +if [[ -n "${JOB_ID}" ]]; then + echo "== GET job: ensure raw api_key is not echoed ==" + curl -fsS "${auth_hdr[@]}" "${BASE_URL}/api/v2/jobs/${JOB_ID}" | tee /tmp/byok_job.json | head -c 800; echo + if grep -q 'sk-smoke-test-key-not-real' /tmp/byok_job.json; then + echo "FAIL: raw API key leaked in job response" >&2 + exit 1 + fi + echo "No raw key in job response OK" +fi + +echo "== v2 jobs: accept llm_config.text + vision (two endpoints) shape ==" +code=$(curl -sS -o /tmp/byok_split.json -w '%{http_code}' "${auth_hdr[@]}" \ + -d '{ + "source_type":"url", + "source_url":"https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", + "file_name":"dummy.pdf", + "data_id":"byok-smoke-split", + "llm_config":{ + "text":{ + "api_key":"sk-smoke-text", + "model":"gpt-4o-mini", + "base_url":"https://api.openai.com/v1" + }, + "vision":{ + "api_key":"sk-smoke-vision", + "model":"qwen-vl-max", + "base_url":"https://dashscope.aliyuncs.com/compatible-mode/v1" + } + } + }' \ + "${BASE_URL}/api/v2/jobs") +echo "HTTP $code" +cat /tmp/byok_split.json | head -c 400; echo +test "$code" = "200" || test "$code" = "201" + +echo "== v2 jobs: accept llm_config.text-only (partial override) shape ==" +code=$(curl -sS -o /tmp/byok_text.json -w '%{http_code}' "${auth_hdr[@]}" \ + -d '{ + "source_type":"url", + "source_url":"https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf", + "file_name":"dummy.pdf", + "data_id":"byok-smoke-text-only", + "llm_config":{ + "text":{ + "api_key":"sk-smoke-test-key-not-real", + "model":"gpt-4o-mini", + "base_url":"https://api.openai.com/v1" + } + } + }' \ + "${BASE_URL}/api/v2/jobs") +echo "HTTP $code" +cat /tmp/byok_text.json | head -c 400; echo +test "$code" = "200" || test "$code" = "201" + +echo "== v1 jobs: llm_config should be rejected as unsupported extra ==" +code=$(curl -sS -o /tmp/byok_v1.json -w '%{http_code}' "${auth_hdr[@]}" \ + -d '{ + "source_type":"url", + "source_url":"https://example.com/a.pdf", + "llm_config":{"text":{"api_key":"sk-x","model":"m","base_url":"https://example.com/v1"}} + }' \ + "${BASE_URL}/api/v1/jobs") +echo "HTTP $code" +cat /tmp/byok_v1.json | head -c 400; echo +# Prefer 422 / validation error; 200 would mean v1 accepted BYOK (bad) +test "$code" != "200" && test "$code" != "201" + +echo +echo "BYOK smoke checks finished."