diff --git a/README.md b/README.md index c55c262..775b524 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,53 @@ print(result.namespace) # "default" or your explicit namespace print(result.document_id) # Published canonical document id ``` +### Bring your own LLM keys (BYOK) + +Pass OpenAI-compatible credentials for parsing or agentic retrieval. + +Flat root applies to both channels (one multimodal model). Use `models` for +different model ids on the same endpoint, or `text` / `vision` for different +provider endpoints: + +```python +# Multimodal shorthand — one model for text + vision +llm_config = { + "api_key": "sk-...", + "model": "gpt-4o", + "base_url": "https://api.openai.com/v1", +} + +# Same endpoint, different models per channel +llm_config = { + "api_key": "sk-...", + "base_url": "https://api.openai.com/v1", + "models": {"text": "gpt-4o-mini", "vision": "gpt-4o"}, +} + +# Or two different endpoints +llm_config = { + "text": { + "api_key": "sk-...", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + "vision": { + "api_key": "sk-ali-...", + "model": "qwen-vl-max", + "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1", + }, +} + +result = client.parse(file=Path("report.pdf"), llm_config=llm_config) + +response = client.retrieval.query( + namespace="support-center", + query="refund policy", + use_agentic=True, + llm_config=llm_config, +) +``` + ### Access different chunk types ```python diff --git a/docs/usage.md b/docs/usage.md index c8956db..6133fd2 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -144,6 +144,7 @@ result = client.parse(file=pdf_bytes, file_name="report.pdf") | `data_id` | `str \| None` | `None` | Your own correlation/idempotency identifier (max 128 chars). | | `parsing_params` | `ParsingParams \| None` | `None` | Parsing configuration (see below). | | `webhook` | `WebhookConfig \| None` | `None` | Webhook for completion notification. | +| `llm_config` | `LLMConfig \| None` | `None` | BYOK OpenAI-compatible credentials (flat root and/or `text` / `vision`). | | `poll_interval` | `float` | `10.0` | Initial polling interval in seconds. | | `poll_timeout` | `float` | `1800.0` | Maximum time to wait for completion (30 min). | | `verify_checksum` | `bool` | `True` | Verify SHA-256 checksum of the downloaded ZIP. | @@ -197,6 +198,24 @@ result = client.parse( ) ``` +### Bring your own LLM keys (BYOK) + +Pass OpenAI-compatible credentials via `llm_config` on `parse()`, `jobs.create()`, +or `retrieval.query()`. Flat root applies to both channels; use `models` for +different model ids on the same endpoint, or `text` / `vision` for different +provider endpoints: + +```python +result = client.parse( + file=Path("report.pdf"), + llm_config={ + "api_key": "sk-...", + "base_url": "https://api.openai.com/v1", + "models": {"text": "gpt-4o-mini", "vision": "gpt-4o"}, + }, +) +``` + --- ## Working with Results @@ -379,6 +398,7 @@ print(result.statistics) | `data_id` | `str \| None` | `None` | Your own correlation/idempotency identifier. | | `parsing_params` | `ParsingParams \| None` | `None` | Parsing configuration. | | `webhook` | `WebhookConfig \| None` | `None` | Webhook for completion notification. | +| `llm_config` | `LLMConfig \| None` | `None` | BYOK OpenAI-compatible credentials (flat root and/or `text` / `vision`). | Returns a `Job` object: @@ -536,6 +556,7 @@ for result in response.results: | `use_agentic` | `bool \| None` | `None` | Force agentic (`True`) or legacy (`False`) retrieval. `None` uses server default. | | `chunk_types` | `list["text" \| "image" \| "table" \| "page"] \| None` | `None` | Restrict retrieval to the selected chunk types. | | `data_type` | `int \| None` | `None` | Deprecated server selector; `7` maps to page chunks and `8` maps to text+image+table. Prefer `chunk_types`. | +| `llm_config` | `LLMConfig \| None` | `None` | BYOK OpenAI-compatible credentials for agentic retrieval (flat root and/or `text` / `vision`). | Retrieval results expose `content`, not the older parse-result `text` field. Media results may include `asset_url` when the server can sign the referenced diff --git a/src/knowhere/__init__.py b/src/knowhere/__init__.py index fa15df5..f338feb 100644 --- a/src/knowhere/__init__.py +++ b/src/knowhere/__init__.py @@ -46,7 +46,13 @@ DocumentListResponse, ) from knowhere.types.job import Job, JobError, JobProgress, JobResult -from knowhere.types.params import ParsingParams, WebhookConfig +from knowhere.types.params import ( + LLMConfig, + LLMModelsConfig, + LLMProviderConfig, + ParsingParams, + WebhookConfig, +) from knowhere.types.retrieval import ( RetrievalChannel, RetrievalChunkType, @@ -145,6 +151,9 @@ "TableChunk", "Chunk", # Param types + "LLMConfig", + "LLMModelsConfig", + "LLMProviderConfig", "ParsingParams", "WebhookConfig", # Callback types diff --git a/src/knowhere/_client.py b/src/knowhere/_client.py index dff40af..dac5b7c 100644 --- a/src/knowhere/_client.py +++ b/src/knowhere/_client.py @@ -23,7 +23,7 @@ from knowhere.resources.jobs import AsyncJobs, Jobs from knowhere.resources.retrieval import AsyncRetrieval, Retrieval from knowhere.types.job import Job, JobResult -from knowhere.types.params import ParsingParams, WebhookConfig +from knowhere.types.params import LLMConfig, ParsingParams, WebhookConfig from knowhere.types.result import ParseResult _logger = getLogger() @@ -66,6 +66,7 @@ def parse( document_id: Optional[str] = ..., parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., + llm_config: Optional[LLMConfig] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -84,6 +85,7 @@ def parse( document_id: Optional[str] = ..., parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., + llm_config: Optional[LLMConfig] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -102,6 +104,7 @@ def parse( document_id: Optional[str] = None, parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, + llm_config: Optional[LLMConfig] = None, poll_interval: float = DEFAULT_POLL_INTERVAL, poll_timeout: float = DEFAULT_POLL_TIMEOUT, verify_checksum: bool = True, @@ -127,6 +130,7 @@ def parse( document_id=document_id, parsing_params=parsing_params, webhook=webhook, + llm_config=llm_config, ) else: resolved_name: Optional[str] = file_name @@ -140,6 +144,7 @@ def parse( document_id=document_id, parsing_params=parsing_params, webhook=webhook, + llm_config=llm_config, ) assert file is not None self.jobs.upload(job, file, on_progress=on_upload_progress) @@ -194,6 +199,7 @@ async def parse( document_id: Optional[str] = ..., parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., + llm_config: Optional[LLMConfig] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -212,6 +218,7 @@ async def parse( document_id: Optional[str] = ..., parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., + llm_config: Optional[LLMConfig] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -230,6 +237,7 @@ async def parse( document_id: Optional[str] = None, parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, + llm_config: Optional[LLMConfig] = None, poll_interval: float = DEFAULT_POLL_INTERVAL, poll_timeout: float = DEFAULT_POLL_TIMEOUT, verify_checksum: bool = True, @@ -251,6 +259,7 @@ async def parse( document_id=document_id, parsing_params=parsing_params, webhook=webhook, + llm_config=llm_config, ) else: resolved_name: Optional[str] = file_name @@ -264,6 +273,7 @@ async def parse( document_id=document_id, parsing_params=parsing_params, webhook=webhook, + llm_config=llm_config, ) assert file is not None await self.jobs.upload(job, file, on_progress=on_upload_progress) diff --git a/src/knowhere/resources/jobs.py b/src/knowhere/resources/jobs.py index faaf061..24e8303 100644 --- a/src/knowhere/resources/jobs.py +++ b/src/knowhere/resources/jobs.py @@ -19,7 +19,7 @@ from knowhere.lib.upload import asyncUploadFile, syncUploadFile from knowhere.resources._base import AsyncAPIResource, SyncAPIResource from knowhere.types.job import Job, JobResult -from knowhere.types.params import ParsingParams, WebhookConfig +from knowhere.types.params import LLMConfig, ParsingParams, WebhookConfig from knowhere.types.result import ParseResult _logger = getLogger() @@ -39,6 +39,7 @@ def create( data_id: Optional[str] = None, parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, + llm_config: Optional[LLMConfig] = None, ) -> Job: """Create a new parsing job. @@ -51,6 +52,7 @@ def create( data_id: Optional idempotency / correlation identifier. parsing_params: Optional parsing configuration. webhook: Optional webhook configuration. + llm_config: Optional BYOK LLM credentials (OpenAI-compatible). Returns: A ``Job`` object with upload details if ``source_type="file"``. @@ -70,6 +72,8 @@ def create( body["parsing_params"] = dict(parsing_params) if webhook is not None: body["webhook"] = dict(webhook) + if llm_config is not None: + body["llm_config"] = dict(llm_config) return self._request( "POST", @@ -195,6 +199,7 @@ async def create( data_id: Optional[str] = None, parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, + llm_config: Optional[LLMConfig] = None, ) -> Job: """Create a new parsing job (async).""" body: Dict[str, Any] = {"source_type": source_type} @@ -212,6 +217,8 @@ async def create( body["parsing_params"] = dict(parsing_params) if webhook is not None: body["webhook"] = dict(webhook) + if llm_config is not None: + body["llm_config"] = dict(llm_config) return await self._request( "POST", diff --git a/src/knowhere/resources/retrieval.py b/src/knowhere/resources/retrieval.py index ad3e940..6061edf 100644 --- a/src/knowhere/resources/retrieval.py +++ b/src/knowhere/resources/retrieval.py @@ -5,6 +5,7 @@ from typing import Any, Dict, Optional from knowhere.resources._base import AsyncAPIResource, SyncAPIResource +from knowhere.types.params import LLMConfig from knowhere.types.retrieval import ( RetrievalChunkType, RetrievalChannel, @@ -35,6 +36,7 @@ def query( internal_recall_k: Optional[int] = None, exclude_document_ids: Optional[list[str]] = None, exclude_sections: Optional[list[RetrievalSectionExclusion]] = None, + llm_config: Optional[LLMConfig] = None, ) -> RetrievalQueryResponse: """Query published documents in a namespace.""" body: Dict[str, Any] = {"query": query} @@ -66,6 +68,8 @@ def query( body["exclude_document_ids"] = exclude_document_ids if exclude_sections is not None: body["exclude_sections"] = exclude_sections + if llm_config is not None: + body["llm_config"] = dict(llm_config) return self._request( "POST", @@ -96,6 +100,7 @@ async def query( internal_recall_k: Optional[int] = None, exclude_document_ids: Optional[list[str]] = None, exclude_sections: Optional[list[RetrievalSectionExclusion]] = None, + llm_config: Optional[LLMConfig] = None, ) -> RetrievalQueryResponse: """Query published documents in a namespace.""" body: Dict[str, Any] = {"query": query} @@ -127,6 +132,8 @@ async def query( body["exclude_document_ids"] = exclude_document_ids if exclude_sections is not None: body["exclude_sections"] = exclude_sections + if llm_config is not None: + body["llm_config"] = dict(llm_config) return await self._request( "POST", diff --git a/src/knowhere/types/__init__.py b/src/knowhere/types/__init__.py index 6588f6f..b767c40 100644 --- a/src/knowhere/types/__init__.py +++ b/src/knowhere/types/__init__.py @@ -13,7 +13,13 @@ DocumentListResponse, ) from knowhere.types.job import Job, JobError, JobResult -from knowhere.types.params import ParsingParams, WebhookConfig +from knowhere.types.params import ( + LLMConfig, + LLMModelsConfig, + LLMProviderConfig, + ParsingParams, + WebhookConfig, +) from knowhere.types.retrieval import ( RetrievalChannel, RetrievalChunkType, @@ -68,6 +74,9 @@ "RetrievalQueryResponse", "RetrievalResult", # params + "LLMConfig", + "LLMModelsConfig", + "LLMProviderConfig", "ParsingParams", "WebhookConfig", # result diff --git a/src/knowhere/types/params.py b/src/knowhere/types/params.py index 8359066..927e700 100644 --- a/src/knowhere/types/params.py +++ b/src/knowhere/types/params.py @@ -19,6 +19,37 @@ class ParsingParams(TypedDict, total=False): add_frag_desc: bool +class LLMProviderConfig(TypedDict, total=False): + """OpenAI-compatible provider credentials for a single modality.""" + + api_key: str + model: str + base_url: str + + +class LLMModelsConfig(TypedDict, total=False): + """Per-channel model ids sharing root api_key / base_url.""" + + text: str + vision: str + + +class LLMConfig(TypedDict, total=False): + """Bring-your-own-key LLM credentials. + + Flat root (``api_key`` / ``model`` / ``base_url``) applies to both channels. + Use ``models`` for different model ids on the same endpoint, or ``text`` / + ``vision`` objects for different provider endpoints. + """ + + api_key: str + model: str + base_url: str + models: LLMModelsConfig + text: LLMProviderConfig + vision: LLMProviderConfig + + class WebhookConfig(TypedDict, total=False): """Webhook configuration for job completion notifications.""" diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 6d09564..c9ace62 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -103,6 +103,75 @@ def test_create_sends_correct_body( assert body["namespace"] == "support-center" assert body["document_id"] == "doc_123" + @respx.mock + def test_create_sends_llm_config( + self, + sync_client: Any, + ) -> None: + """POST body includes nested llm_config when provided.""" + response_body: Dict[str, Any] = { + "job_id": "job_llm_config", + "status": "pending", + "source_type": "url", + } + + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=response_body)) + + sync_client.jobs.create( + source_type="url", + source_url="https://example.com/doc.pdf", + llm_config={ + "text": { + "api_key": "sk-text", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + "vision": { + "api_key": "sk-vision", + "model": "gpt-4o", + "base_url": "https://api.openai.com/v1", + }, + }, + ) + + assert route.called + body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert body["llm_config"] == { + "text": { + "api_key": "sk-text", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + "vision": { + "api_key": "sk-vision", + "model": "gpt-4o", + "base_url": "https://api.openai.com/v1", + }, + } + + @respx.mock + def test_create_omits_llm_config_when_none( + self, + sync_client: Any, + ) -> None: + """llm_config is omitted from the POST body when not set.""" + response_body: Dict[str, Any] = { + "job_id": "job_no_llm", + "status": "pending", + "source_type": "url", + } + + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=response_body)) + + sync_client.jobs.create( + source_type="url", + source_url="https://example.com/doc.pdf", + ) + + assert route.called + body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert "llm_config" not in body + # --------------------------------------------------------------------------- # jobs.get() diff --git a/tests/test_parse.py b/tests/test_parse.py index f12d6b2..3559797 100644 --- a/tests/test_parse.py +++ b/tests/test_parse.py @@ -101,6 +101,56 @@ def test_parse_url_full_flow( assert parse_result.namespace == "support-center" assert parse_result.document_id == "doc_123" + @respx.mock + def test_parse_url_forwards_llm_config( + self, + sync_client: Any, + sample_zip_bytes: bytes, + ) -> None: + """parse() threads llm_config into jobs.create POST body.""" + import json + + job_id: str = "job_llm_parse" + result_url: str = "https://storage.example.com/result.zip" + llm_config = { + "text": { + "api_key": "sk-text", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + } + + create_route = respx.post(JOBS_URL).mock( + return_value=httpx.Response( + 200, + json=_make_create_response(job_id, "url"), + ) + ) + respx.get(f"{JOBS_URL}/{job_id}").mock( + return_value=httpx.Response( + 200, + json=_make_done_response(job_id, result_url), + ) + ) + respx.get(result_url).mock( + return_value=httpx.Response( + 200, + content=sample_zip_bytes, + headers={"Content-Type": "application/zip"}, + ) + ) + + sync_client.parse( + url="https://example.com/doc.pdf", + llm_config=llm_config, + poll_interval=0.01, + verify_checksum=False, + ) + + assert create_route.called + body: Dict[str, Any] = json.loads(create_route.calls[0].request.read()) + assert body["llm_config"] == llm_config + # --------------------------------------------------------------------------- # parse(file=Path(...)) diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py index 3632674..6529626 100644 --- a/tests/test_retrieval.py +++ b/tests/test_retrieval.py @@ -179,6 +179,56 @@ def test_query_sends_chunk_types( } assert response.results[0].chunk_id == "chunk_001" + @respx.mock + def test_query_sends_llm_config(self, sync_client: Any) -> None: + """llm_config with nested text/vision is included in the POST body.""" + route = respx.post(RETRIEVAL_QUERY_URL).mock( + return_value=httpx.Response(200, json=_make_retrieval_response()) + ) + + sync_client.retrieval.query( + query="refund policy", + llm_config={ + "text": { + "api_key": "sk-text", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + "vision": { + "api_key": "sk-vision", + "model": "gpt-4o", + "base_url": "https://api.openai.com/v1", + }, + }, + ) + + assert route.called + request_body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert request_body["llm_config"] == { + "text": { + "api_key": "sk-text", + "model": "gpt-4o-mini", + "base_url": "https://api.openai.com/v1", + }, + "vision": { + "api_key": "sk-vision", + "model": "gpt-4o", + "base_url": "https://api.openai.com/v1", + }, + } + + @respx.mock + def test_query_omits_llm_config_when_none(self, sync_client: Any) -> None: + """llm_config is omitted from the POST body when not set.""" + route = respx.post(RETRIEVAL_QUERY_URL).mock( + return_value=httpx.Response(200, json=_make_retrieval_response()) + ) + + sync_client.retrieval.query(query="refund policy") + + request_body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert "llm_config" not in request_body + @respx.mock def test_query_omits_defaulted_optional_fields(self, sync_client: Any) -> None: route = respx.post(RETRIEVAL_QUERY_URL).mock(