diff --git a/README.md b/README.md index 775b524..440bd83 100644 --- a/README.md +++ b/README.md @@ -42,19 +42,20 @@ for page in result.page_chunks: ## Retrieval and document lifecycle New documents are published into a retrieval namespace. The server returns a -stable `document_id` after the job is published. `client.jobs.create(...)` -does not return a usable `document_id`; persist `job_result.document_id` if you -need to update or archive the same document later. +stable `document_id` on job create when it has a planned id, and on the +completed `job_result` after publication. ```python job = client.jobs.create( source_type="url", source_url="https://example.com/manual.pdf", namespace="support-center", + document_metadata={"title": "Support manual"}, ) +document_id = job.document_id job_result = client.jobs.wait(job.job_id) -document_id = job_result.document_id +document_id = document_id or job_result.document_id if document_id is None: raise RuntimeError("Expected document_id after successful publication.") @@ -119,6 +120,11 @@ if chunks.chunks: print(chunk.chunk.content) print(chunk.chunk.metadata.get("page_nums")) # Page citations. print(chunk.chunk.asset_url) # Requested 7-day URL when available. + page_assets = chunk.chunk.metadata.get("pageAssets") or [] + print(page_assets) + +source = client.documents.get_page_citation_source(document_id) +print(source.url) client.documents.archive(document_id) ``` @@ -152,6 +158,12 @@ response = client.retrieval.query( While you can provide an `api_key` keyword argument, we recommend using [python-dotenv](https://pypi.org/project/python-dotenv/) to add `KNOWHERE_API_KEY="sk_..."` to your `.env` file so that your API key is not stored in source control. +Short-lived dashboard tokens can use `auth_token_provider` instead of a static key. If `api_key` is also set, the static key wins. + +```python +client = knowhere.Knowhere(auth_token_provider=lambda: current_access_token()) +``` + ### Parse a local file ```python diff --git a/docs/usage.md b/docs/usage.md index 6133fd2..8b60424 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -42,7 +42,7 @@ uv add knowhere-python-sdk ## Authentication -The SDK requires an API key. You can provide it in three ways (highest priority first): +Provide a static API key or a short-lived token provider. 1. Constructor argument: @@ -76,6 +76,12 @@ load_dotenv() client = knowhere.Knowhere() ``` +4. Short-lived bearer token. If `api_key` is also set (or `KNOWHERE_API_KEY` is present), the static key wins. + +```python +client = knowhere.Knowhere(auth_token_provider=lambda: current_access_token()) +``` + ## Quick Start ```python @@ -145,6 +151,7 @@ result = client.parse(file=pdf_bytes, file_name="report.pdf") | `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`). | +| `document_metadata` | `dict \| None` | `None` | Display metadata copied onto the published document. Official `created_by_client` / `client_version` defaults are filled when omitted. | | `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. | @@ -399,6 +406,7 @@ print(result.statistics) | `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`). | +| `document_metadata` | `dict \| None` | `None` | Display metadata copied onto the published document. Official `created_by_client` / `client_version` defaults are filled when omitted. | Returns a `Job` object: @@ -627,6 +635,9 @@ image_chunk = client.documents.get_chunk( ) print(image_chunk.chunk.asset_url) +source = client.documents.get_page_citation_source("doc_123") +print(source.url, source.content_type) + archived = client.documents.archive("doc_123") print(archived.status) # "archived" ``` diff --git a/pyproject.toml b/pyproject.toml index e3d1abf..0b3b4e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dev = [ "pytest>=7.0.0,<9.0.0", "pytest-asyncio>=0.23.0", "respx>=0.21.0", - "ruff>=0.1.0", + "ruff>=0.15.0,<0.16.0", "mypy>=1.0.0", "coverage>=7.0.0", ] diff --git a/src/knowhere/__init__.py b/src/knowhere/__init__.py index f338feb..f865051 100644 --- a/src/knowhere/__init__.py +++ b/src/knowhere/__init__.py @@ -33,8 +33,12 @@ ServiceUnavailableError, ValidationError, ) -from knowhere._types import PollProgressCallback, UploadProgressCallback +from knowhere._types import AuthTokenProvider, PollProgressCallback, UploadProgressCallback from knowhere._version import __version__ +from knowhere.lib.document_metadata import ( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + merge_document_metadata_defaults, +) from knowhere.types.document import ( Document, DocumentChunk, @@ -44,25 +48,23 @@ DocumentChunkType, DocumentListPagination, DocumentListResponse, + DocumentPageCitationSource, ) from knowhere.types.job import Job, JobError, JobProgress, JobResult +from knowhere.types.page_citation import ( + PAGE_CITATION_ASSETS_METADATA_KEY, + PageCitationAsset, + PageCitationAssetContentType, + PageCitationAssetSource, +) from knowhere.types.params import ( + DocumentMetadata, LLMConfig, LLMModelsConfig, LLMProviderConfig, ParsingParams, WebhookConfig, ) -from knowhere.types.retrieval import ( - RetrievalChannel, - RetrievalChunkType, - RetrievalFilterMode, - RetrievalReferencedChunk, - RetrievalSectionExclusion, - RetrievalSource, - RetrievalQueryResponse, - RetrievalResult, -) from knowhere.types.result import ( BaseChunk, Checksum, @@ -82,6 +84,16 @@ TableFileInfo, TextChunk, ) +from knowhere.types.retrieval import ( + RetrievalChannel, + RetrievalChunkType, + RetrievalFilterMode, + RetrievalQueryResponse, + RetrievalReferencedChunk, + RetrievalResult, + RetrievalSectionExclusion, + RetrievalSource, +) __all__: list[str] = [ # Clients @@ -89,6 +101,8 @@ "AsyncKnowhere", # Version "__version__", + "PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS", + "merge_document_metadata_defaults", # Exceptions "KnowhereError", "ValidationError", @@ -123,6 +137,11 @@ "DocumentChunkType", "DocumentListPagination", "DocumentListResponse", + "DocumentPageCitationSource", + "PageCitationAsset", + "PageCitationAssetContentType", + "PageCitationAssetSource", + "PAGE_CITATION_ASSETS_METADATA_KEY", # Retrieval types "RetrievalChannel", "RetrievalChunkType", @@ -151,6 +170,7 @@ "TableChunk", "Chunk", # Param types + "DocumentMetadata", "LLMConfig", "LLMModelsConfig", "LLMProviderConfig", @@ -159,4 +179,5 @@ # Callback types "UploadProgressCallback", "PollProgressCallback", + "AuthTokenProvider", ] diff --git a/src/knowhere/_base_client.py b/src/knowhere/_base_client.py index 5994ff9..b9aec9e 100644 --- a/src/knowhere/_base_client.py +++ b/src/knowhere/_base_client.py @@ -6,6 +6,7 @@ from __future__ import annotations +import inspect import os import random import time @@ -30,6 +31,7 @@ ) from knowhere._logging import getLogger, redactSensitiveHeaders from knowhere._response import APIResponse +from knowhere._types import AuthTokenProvider from knowhere._version import __version__ T = TypeVar("T") @@ -60,7 +62,8 @@ class BaseClient: """Shared configuration and helper methods for sync/async clients.""" - api_key: str + api_key: Optional[str] + _auth_token_provider: Optional[AuthTokenProvider] base_url: str timeout: float upload_timeout: float @@ -71,20 +74,26 @@ def __init__( self, *, api_key: Optional[str] = None, + auth_token_provider: Optional[AuthTokenProvider] = None, base_url: Optional[str] = None, timeout: Optional[float] = None, upload_timeout: Optional[float] = None, max_retries: Optional[int] = None, default_headers: Optional[Dict[str, str]] = None, ) -> None: - # Resolve: arg > env > default + # Resolve: arg > env > default. A static api_key wins over the provider. resolved_key: Optional[str] = api_key or os.environ.get(ENV_API_KEY) - if not resolved_key: + if resolved_key: + self.api_key = resolved_key + self._auth_token_provider = None + elif auth_token_provider is not None: + self.api_key = None + self._auth_token_provider = auth_token_provider + else: raise ValidationError( - "An API key must be provided via the 'api_key' argument " - f"or the {ENV_API_KEY} environment variable." + "An API key must be provided via the 'api_key' argument, " + f"the {ENV_API_KEY} environment variable, or auth_token_provider." ) - self.api_key = resolved_key self.base_url = (base_url or os.environ.get(ENV_BASE_URL) or DEFAULT_BASE_URL).rstrip("/") self.timeout = timeout if timeout is not None else DEFAULT_TIMEOUT self.upload_timeout = ( @@ -93,10 +102,38 @@ def __init__( self.max_retries = max_retries if max_retries is not None else DEFAULT_MAX_RETRIES self._default_headers = default_headers or {} - def _buildHeaders(self) -> Dict[str, str]: + def _sync_auth_token(self) -> str: + """Resolve a bearer token for a synchronous request.""" + if self.api_key: + return self.api_key + if self._auth_token_provider is None: + raise ValidationError("Authentication token provider is not configured.") + token = self._auth_token_provider() + if inspect.isawaitable(token): + raise ValidationError( + "Synchronous clients require a non-async auth_token_provider." + ) + if not token: + raise ValidationError("Authentication token provider returned an empty token") + return str(token) + + async def _async_auth_token(self) -> str: + """Resolve a bearer token for an asynchronous request.""" + if self.api_key: + return self.api_key + if self._auth_token_provider is None: + raise ValidationError("Authentication token provider is not configured.") + token = self._auth_token_provider() + if inspect.isawaitable(token): + token = await token + if not token: + raise ValidationError("Authentication token provider returned an empty token") + return str(token) + + def _buildHeaders(self, *, token: str) -> Dict[str, str]: """Return headers including auth and user-agent.""" headers: Dict[str, str] = { - "Authorization": f"Bearer {self.api_key}", + "Authorization": f"Bearer {token}", "User-Agent": f"knowhere-python/{__version__}", "Accept": "application/json", } @@ -214,6 +251,7 @@ def __init__( self, *, api_key: Optional[str] = None, + auth_token_provider: Optional[AuthTokenProvider] = None, base_url: Optional[str] = None, timeout: Optional[float] = None, upload_timeout: Optional[float] = None, @@ -222,6 +260,7 @@ def __init__( ) -> None: super().__init__( api_key=api_key, + auth_token_provider=auth_token_provider, base_url=base_url, timeout=timeout, upload_timeout=upload_timeout, @@ -248,20 +287,19 @@ def _request( ) -> T: """Execute an HTTP request with automatic retries and error handling.""" url: str = self._buildRequestUrl(path) - request_headers: Dict[str, str] = self._buildHeaders() - if headers: - request_headers.update(headers) - effective_timeout: float = timeout if timeout is not None else self.timeout - _logger.debug( - "Request: %s %s headers=%s", - method, - url, - redactSensitiveHeaders(request_headers), - ) - for attempt in range(self.max_retries + 1): + request_headers: Dict[str, str] = self._buildHeaders(token=self._sync_auth_token()) + if headers: + request_headers.update(headers) + + _logger.debug( + "Request: %s %s headers=%s", + method, + url, + redactSensitiveHeaders(request_headers), + ) try: response: httpx.Response = self._client.request( method, @@ -362,6 +400,7 @@ def __init__( self, *, api_key: Optional[str] = None, + auth_token_provider: Optional[AuthTokenProvider] = None, base_url: Optional[str] = None, timeout: Optional[float] = None, upload_timeout: Optional[float] = None, @@ -370,6 +409,7 @@ def __init__( ) -> None: super().__init__( api_key=api_key, + auth_token_provider=auth_token_provider, base_url=base_url, timeout=timeout, upload_timeout=upload_timeout, @@ -396,20 +436,21 @@ async def _request( import asyncio url: str = self._buildRequestUrl(path) - request_headers: Dict[str, str] = self._buildHeaders() - if headers: - request_headers.update(headers) - effective_timeout: float = timeout if timeout is not None else self.timeout - _logger.debug( - "Async request: %s %s headers=%s", - method, - url, - redactSensitiveHeaders(request_headers), - ) - for attempt in range(self.max_retries + 1): + request_headers: Dict[str, str] = self._buildHeaders( + token=await self._async_auth_token() + ) + if headers: + request_headers.update(headers) + + _logger.debug( + "Async request: %s %s headers=%s", + method, + url, + redactSensitiveHeaders(request_headers), + ) try: response: httpx.Response = await self._client.request( method, diff --git a/src/knowhere/_client.py b/src/knowhere/_client.py index dac5b7c..96fc660 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 LLMConfig, ParsingParams, WebhookConfig +from knowhere.types.params import DocumentMetadata, LLMConfig, ParsingParams, WebhookConfig from knowhere.types.result import ParseResult _logger = getLogger() @@ -67,6 +67,7 @@ def parse( parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., llm_config: Optional[LLMConfig] = ..., + document_metadata: Optional[DocumentMetadata] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -86,6 +87,7 @@ def parse( parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., llm_config: Optional[LLMConfig] = ..., + document_metadata: Optional[DocumentMetadata] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -105,6 +107,7 @@ def parse( parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, llm_config: Optional[LLMConfig] = None, + document_metadata: Optional[DocumentMetadata] = None, poll_interval: float = DEFAULT_POLL_INTERVAL, poll_timeout: float = DEFAULT_POLL_TIMEOUT, verify_checksum: bool = True, @@ -131,6 +134,7 @@ def parse( parsing_params=parsing_params, webhook=webhook, llm_config=llm_config, + document_metadata=document_metadata, ) else: resolved_name: Optional[str] = file_name @@ -145,6 +149,7 @@ def parse( parsing_params=parsing_params, webhook=webhook, llm_config=llm_config, + document_metadata=document_metadata, ) assert file is not None self.jobs.upload(job, file, on_progress=on_upload_progress) @@ -200,6 +205,7 @@ async def parse( parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., llm_config: Optional[LLMConfig] = ..., + document_metadata: Optional[DocumentMetadata] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -219,6 +225,7 @@ async def parse( parsing_params: Optional[ParsingParams] = ..., webhook: Optional[WebhookConfig] = ..., llm_config: Optional[LLMConfig] = ..., + document_metadata: Optional[DocumentMetadata] = ..., poll_interval: float = ..., poll_timeout: float = ..., verify_checksum: bool = ..., @@ -238,6 +245,7 @@ async def parse( parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, llm_config: Optional[LLMConfig] = None, + document_metadata: Optional[DocumentMetadata] = None, poll_interval: float = DEFAULT_POLL_INTERVAL, poll_timeout: float = DEFAULT_POLL_TIMEOUT, verify_checksum: bool = True, @@ -260,6 +268,7 @@ async def parse( parsing_params=parsing_params, webhook=webhook, llm_config=llm_config, + document_metadata=document_metadata, ) else: resolved_name: Optional[str] = file_name @@ -274,6 +283,7 @@ async def parse( parsing_params=parsing_params, webhook=webhook, llm_config=llm_config, + document_metadata=document_metadata, ) assert file is not None await self.jobs.upload(job, file, on_progress=on_upload_progress) diff --git a/src/knowhere/_types.py b/src/knowhere/_types.py index ffcb787..844dc35 100644 --- a/src/knowhere/_types.py +++ b/src/knowhere/_types.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Optional, Union +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Union from typing_extensions import TypeAlias @@ -54,3 +54,6 @@ def __bool__(self) -> bool: # Poll progress: (current_job_result, elapsed_seconds) # We use a string forward-ref to avoid a circular import with types.job PollProgressCallback: TypeAlias = Callable[["JobResult", float], None] + +# Bearer token provider: sync str, or async str for AsyncKnowhere +AuthTokenProvider: TypeAlias = Callable[[], Union[str, Awaitable[str]]] diff --git a/src/knowhere/lib/document_metadata.py b/src/knowhere/lib/document_metadata.py new file mode 100644 index 0000000..adef237 --- /dev/null +++ b/src/knowhere/lib/document_metadata.py @@ -0,0 +1,27 @@ +"""Official-client document metadata defaults for job creates.""" + +from __future__ import annotations + +from typing import Any, Dict, Mapping, Optional + +from knowhere._version import __version__ + +# Wire format: ``{ created_by_client, client_version }``. +PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS: Dict[str, Any] = { + "created_by_client": "python-sdk", + "client_version": __version__, +} + + +def merge_document_metadata_defaults( + defaults: Mapping[str, Any], + provided: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Merge official-client defaults under caller-provided metadata. + + Caller keys win; defaults fill keys that are missing. + """ + merged: Dict[str, Any] = dict(defaults) + if provided: + merged.update(dict(provided)) + return merged diff --git a/src/knowhere/resources/documents.py b/src/knowhere/resources/documents.py index e1cc183..c69f653 100644 --- a/src/knowhere/resources/documents.py +++ b/src/knowhere/resources/documents.py @@ -11,6 +11,7 @@ DocumentChunkResponse, DocumentChunkType, DocumentListResponse, + DocumentPageCitationSource, ) @@ -89,6 +90,14 @@ def get_chunk( cast_to=DocumentChunkResponse, ) + def get_page_citation_source(self, document_id: str) -> DocumentPageCitationSource: + """Get a signed normalized source-file URL for citations and debugging.""" + return self._request( + "GET", + self._versionedPath(f"documents/{document_id}/files/page-citation-source"), + cast_to=DocumentPageCitationSource, + ) + def archive(self, document_id: str) -> Document: """Archive one canonical document by ID.""" return self._request( @@ -173,6 +182,14 @@ async def get_chunk( cast_to=DocumentChunkResponse, ) + async def get_page_citation_source(self, document_id: str) -> DocumentPageCitationSource: + """Get a signed normalized source-file URL for citations and debugging.""" + return await self._request( + "GET", + self._versionedPath(f"documents/{document_id}/files/page-citation-source"), + cast_to=DocumentPageCitationSource, + ) + async def archive(self, document_id: str) -> Document: """Archive one canonical document by ID.""" return await self._request( diff --git a/src/knowhere/resources/jobs.py b/src/knowhere/resources/jobs.py index 24e8303..9682b1a 100644 --- a/src/knowhere/resources/jobs.py +++ b/src/knowhere/resources/jobs.py @@ -14,17 +14,61 @@ PollProgressCallback, UploadProgressCallback, ) +from knowhere.lib.document_metadata import ( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + merge_document_metadata_defaults, +) from knowhere.lib.polling import asyncPoll, syncPoll from knowhere.lib.result_parser import parseResultZip 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 LLMConfig, ParsingParams, WebhookConfig +from knowhere.types.params import DocumentMetadata, LLMConfig, ParsingParams, WebhookConfig from knowhere.types.result import ParseResult _logger = getLogger() +def _build_job_create_body( + *, + source_type: str, + source_url: Optional[str], + file_name: Optional[str], + namespace: Optional[str], + document_id: Optional[str], + data_id: Optional[str], + parsing_params: Optional[ParsingParams], + webhook: Optional[WebhookConfig], + llm_config: Optional[LLMConfig], + document_metadata: Optional[DocumentMetadata], +) -> Dict[str, Any]: + """Build a ``POST /v2/jobs`` body, always attaching telemetry metadata.""" + body: Dict[str, Any] = { + "source_type": source_type, + "document_metadata": merge_document_metadata_defaults( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + document_metadata, + ), + } + if source_url is not None: + body["source_url"] = source_url + if file_name is not None: + body["file_name"] = file_name + if namespace is not None: + body["namespace"] = namespace + if document_id is not None: + body["document_id"] = document_id + if data_id is not None: + body["data_id"] = data_id + if parsing_params is not None: + 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 body + + class Jobs(SyncAPIResource): """Synchronous interface for the ``/v2/jobs`` endpoints.""" @@ -40,6 +84,7 @@ def create( parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, llm_config: Optional[LLMConfig] = None, + document_metadata: Optional[DocumentMetadata] = None, ) -> Job: """Create a new parsing job. @@ -53,32 +98,28 @@ def create( parsing_params: Optional parsing configuration. webhook: Optional webhook configuration. llm_config: Optional BYOK LLM credentials (OpenAI-compatible). + document_metadata: Display metadata copied onto the published + document. Official ``created_by_client`` / ``client_version`` + defaults are filled when omitted; caller keys win. Returns: A ``Job`` object with upload details if ``source_type="file"``. """ - body: Dict[str, Any] = {"source_type": source_type} - if source_url is not None: - body["source_url"] = source_url - if file_name is not None: - body["file_name"] = file_name - if namespace is not None: - body["namespace"] = namespace - if document_id is not None: - body["document_id"] = document_id - if data_id is not None: - body["data_id"] = data_id - if parsing_params is not None: - 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", self._versionedPath("jobs"), - body=body, + body=_build_job_create_body( + source_type=source_type, + source_url=source_url, + file_name=file_name, + namespace=namespace, + document_id=document_id, + data_id=data_id, + parsing_params=parsing_params, + webhook=webhook, + llm_config=llm_config, + document_metadata=document_metadata, + ), cast_to=Job, ) @@ -200,30 +241,24 @@ async def create( parsing_params: Optional[ParsingParams] = None, webhook: Optional[WebhookConfig] = None, llm_config: Optional[LLMConfig] = None, + document_metadata: Optional[DocumentMetadata] = None, ) -> Job: """Create a new parsing job (async).""" - body: Dict[str, Any] = {"source_type": source_type} - if source_url is not None: - body["source_url"] = source_url - if file_name is not None: - body["file_name"] = file_name - if namespace is not None: - body["namespace"] = namespace - if document_id is not None: - body["document_id"] = document_id - if data_id is not None: - body["data_id"] = data_id - if parsing_params is not None: - 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", self._versionedPath("jobs"), - body=body, + body=_build_job_create_body( + source_type=source_type, + source_url=source_url, + file_name=file_name, + namespace=namespace, + document_id=document_id, + data_id=data_id, + parsing_params=parsing_params, + webhook=webhook, + llm_config=llm_config, + document_metadata=document_metadata, + ), cast_to=Job, ) diff --git a/src/knowhere/types/__init__.py b/src/knowhere/types/__init__.py index b767c40..7dba489 100644 --- a/src/knowhere/types/__init__.py +++ b/src/knowhere/types/__init__.py @@ -11,25 +11,23 @@ DocumentChunkType, DocumentListPagination, DocumentListResponse, + DocumentPageCitationSource, ) from knowhere.types.job import Job, JobError, JobResult +from knowhere.types.page_citation import ( + PAGE_CITATION_ASSETS_METADATA_KEY, + PageCitationAsset, + PageCitationAssetContentType, + PageCitationAssetSource, +) from knowhere.types.params import ( + DocumentMetadata, LLMConfig, LLMModelsConfig, LLMProviderConfig, ParsingParams, WebhookConfig, ) -from knowhere.types.retrieval import ( - RetrievalChannel, - RetrievalChunkType, - RetrievalFilterMode, - RetrievalReferencedChunk, - RetrievalSectionExclusion, - RetrievalSource, - RetrievalQueryResponse, - RetrievalResult, -) from knowhere.types.result import ( BaseChunk, Checksum, @@ -49,6 +47,16 @@ TableFileInfo, TextChunk, ) +from knowhere.types.retrieval import ( + RetrievalChannel, + RetrievalChunkType, + RetrievalFilterMode, + RetrievalQueryResponse, + RetrievalReferencedChunk, + RetrievalResult, + RetrievalSectionExclusion, + RetrievalSource, +) __all__: list[str] = [ # job @@ -64,6 +72,11 @@ "DocumentChunkType", "DocumentListPagination", "DocumentListResponse", + "DocumentPageCitationSource", + "PageCitationAsset", + "PageCitationAssetContentType", + "PageCitationAssetSource", + "PAGE_CITATION_ASSETS_METADATA_KEY", # retrieval "RetrievalChannel", "RetrievalChunkType", @@ -74,6 +87,7 @@ "RetrievalQueryResponse", "RetrievalResult", # params + "DocumentMetadata", "LLMConfig", "LLMModelsConfig", "LLMProviderConfig", diff --git a/src/knowhere/types/document.py b/src/knowhere/types/document.py index 577f03e..1ec136c 100644 --- a/src/knowhere/types/document.py +++ b/src/knowhere/types/document.py @@ -16,6 +16,7 @@ class Document(BaseModel): status: str current_job_result_id: Optional[str] = None source_file_name: Optional[str] = None + document_metadata: Optional[Dict[str, Any]] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None archived_at: Optional[datetime] = None @@ -104,3 +105,17 @@ class DocumentChunkResponse(BaseModel): job_result_id: Optional[str] = None job_id: Optional[str] = None chunk: DocumentChunk + + +class DocumentPageCitationSource(BaseModel): + """Response from ``GET /v2/documents/{document_id}/files/page-citation-source``.""" + + document_id: str + namespace: Optional[str] = None + job_id: Optional[str] = None + job_result_id: Optional[str] = None + variant: Optional[str] = None + file_name: str + content_type: str + url: str + expires_at: Optional[datetime] = None diff --git a/src/knowhere/types/job.py b/src/knowhere/types/job.py index 537a562..2f1bb15 100644 --- a/src/knowhere/types/job.py +++ b/src/knowhere/types/job.py @@ -41,6 +41,7 @@ class Job(BaseModel): status: str source_type: str namespace: Optional[str] = None + document_id: Optional[str] = None data_id: Optional[str] = None created_at: Optional[datetime] = None upload_url: Optional[str] = None diff --git a/src/knowhere/types/page_citation.py b/src/knowhere/types/page_citation.py new file mode 100644 index 0000000..f27a5e3 --- /dev/null +++ b/src/knowhere/types/page_citation.py @@ -0,0 +1,28 @@ +"""Typed page-citation asset descriptors stored on chunk metadata.""" + +from __future__ import annotations + +from typing import Literal, Optional + +from pydantic import BaseModel + +PAGE_CITATION_ASSETS_METADATA_KEY = "pageAssets" + +PageCitationAssetContentType = Literal["image/png", "image/jpeg"] +PageCitationAssetSource = Literal["knowhere-rendered-page-citation-source"] + + +class PageCitationAsset(BaseModel): + """Server-provided page citation asset descriptor. + + Stored on chunk metadata under ``pageAssets``. The SDK does not generate + these assets; it only types descriptors returned by Knowhere. + """ + + page_num: int + artifact_ref: str + asset_url: Optional[str] = None + content_type: PageCitationAssetContentType + width: Optional[int] = None + height: Optional[int] = None + source: PageCitationAssetSource diff --git a/src/knowhere/types/params.py b/src/knowhere/types/params.py index 927e700..485bf23 100644 --- a/src/knowhere/types/params.py +++ b/src/knowhere/types/params.py @@ -2,8 +2,13 @@ from __future__ import annotations +from typing import Any, Dict + from typing_extensions import TypedDict +DocumentMetadata = Dict[str, Any] +"""Client-provided display metadata copied onto the published document.""" + class ParsingParams(TypedDict, total=False): """Optional parsing parameters for job creation.""" diff --git a/tests/test_client.py b/tests/test_client.py index 2bef3e4..b5a3476 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -6,7 +6,9 @@ from typing import Any from unittest.mock import patch +import httpx import pytest +import respx from knowhere._exceptions import ValidationError @@ -17,6 +19,8 @@ DEFAULT_UPLOAD_TIMEOUT, ) +from tests.conftest import BASE_URL + # --------------------------------------------------------------------------- # Sync client: Knowhere @@ -130,9 +134,69 @@ def test_documents_property_returns_documents_instance(self) -> None: documents: Any = client.documents assert hasattr(documents, "list") assert hasattr(documents, "get") + assert hasattr(documents, "get_page_citation_source") assert hasattr(documents, "archive") client.close() + def test_constructor_accepts_auth_token_provider(self) -> None: + from knowhere import Knowhere + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("KNOWHERE_API_KEY", None) + client: Knowhere = Knowhere(auth_token_provider=lambda: "jwt_test") + assert client.api_key is None + client.close() + + def test_api_key_wins_over_auth_token_provider(self) -> None: + from knowhere import Knowhere + + client: Knowhere = Knowhere( + api_key="sk_static", + auth_token_provider=lambda: "jwt_ignored", + ) + assert client.api_key == "sk_static" + client.close() + + @respx.mock + def test_auth_token_provider_is_called_per_request(self) -> None: + from knowhere import Knowhere + + tokens = iter(["jwt_one", "jwt_two"]) + route = respx.get(f"{BASE_URL}/v2/jobs/job_1").mock( + return_value=httpx.Response( + 200, + json={"job_id": "job_1", "status": "done", "source_type": "url"}, + ) + ) + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("KNOWHERE_API_KEY", None) + client = Knowhere( + auth_token_provider=lambda: next(tokens), + base_url=BASE_URL, + ) + try: + client.jobs.get("job_1") + client.jobs.get("job_1") + finally: + client.close() + + assert route.call_count == 2 + assert route.calls[0].request.headers["Authorization"] == "Bearer jwt_one" + assert route.calls[1].request.headers["Authorization"] == "Bearer jwt_two" + + def test_empty_auth_token_provider_raises_validation_error(self) -> None: + from knowhere import Knowhere + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("KNOWHERE_API_KEY", None) + client = Knowhere(auth_token_provider=lambda: "", base_url=BASE_URL) + try: + with pytest.raises(ValidationError, match="empty token"): + client.jobs.get("job_1") + finally: + client.close() + def test_base_url_trailing_slash_stripped(self) -> None: from knowhere import Knowhere @@ -233,4 +297,34 @@ def test_documents_property_returns_async_documents_instance(self) -> None: documents: Any = client.documents assert hasattr(documents, "list") assert hasattr(documents, "get") + assert hasattr(documents, "get_page_citation_source") assert hasattr(documents, "archive") + + @respx.mock + @pytest.mark.asyncio + async def test_async_auth_token_provider(self) -> None: + from knowhere import AsyncKnowhere + + async def provide_token() -> str: + return "jwt_async" + + route = respx.get(f"{BASE_URL}/v2/jobs/job_1").mock( + return_value=httpx.Response( + 200, + json={"job_id": "job_1", "status": "done", "source_type": "url"}, + ) + ) + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("KNOWHERE_API_KEY", None) + client = AsyncKnowhere( + auth_token_provider=provide_token, + base_url=BASE_URL, + ) + try: + await client.jobs.get("job_1") + finally: + await client.close() + + assert route.called + assert route.calls[0].request.headers["Authorization"] == "Bearer jwt_async" diff --git a/tests/test_document_metadata.py b/tests/test_document_metadata.py new file mode 100644 index 0000000..38bf768 --- /dev/null +++ b/tests/test_document_metadata.py @@ -0,0 +1,41 @@ +"""Tests for official-client document metadata merge.""" + +from __future__ import annotations + +from knowhere._version import __version__ +from knowhere.lib.document_metadata import ( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + merge_document_metadata_defaults, +) + + +class TestMergeDocumentMetadataDefaults: + def test_fills_defaults_when_metadata_omitted(self) -> None: + assert merge_document_metadata_defaults(PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS) == { + "created_by_client": "python-sdk", + "client_version": __version__, + } + + def test_fills_only_missing_keys(self) -> None: + assert merge_document_metadata_defaults( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + {"title": "Report.pdf"}, + ) == { + "created_by_client": "python-sdk", + "client_version": __version__, + "title": "Report.pdf", + } + + def test_caller_keys_win(self) -> None: + assert merge_document_metadata_defaults( + PYTHON_SDK_DOCUMENT_METADATA_DEFAULTS, + { + "created_by_client": "cli", + "client_version": "9.9.9", + "title": "Report.pdf", + }, + ) == { + "created_by_client": "cli", + "client_version": "9.9.9", + "title": "Report.pdf", + } diff --git a/tests/test_documents.py b/tests/test_documents.py index 94f71a8..230559d 100644 --- a/tests/test_documents.py +++ b/tests/test_documents.py @@ -21,6 +21,7 @@ def _make_document(status: str = "active") -> Dict[str, Any]: "status": status, "current_job_result_id": "result_123", "source_file_name": "refund-policy.md", + "document_metadata": {"created_by_client": "python-sdk", "title": "Refund policy"}, "created_at": "2026-04-21T08:00:00Z", "updated_at": "2026-04-21T08:30:00Z", "archived_at": "2026-04-21T09:00:00Z" if status == "archived" else None, @@ -137,6 +138,10 @@ def test_get_document_returns_document_state(self, sync_client: Any) -> None: assert route.called assert document.document_id == "doc_123" assert document.status == "active" + assert document.document_metadata == { + "created_by_client": "python-sdk", + "title": "Refund policy", + } @respx.mock def test_list_chunks_sends_optional_query_params(self, sync_client: Any) -> None: @@ -308,6 +313,95 @@ def test_archive_document_returns_archived_state(self, sync_client: Any) -> None assert document.status == "archived" assert document.archived_at is not None + @respx.mock + def test_get_page_citation_source_hits_canonical_route(self, sync_client: Any) -> None: + route = respx.get(f"{DOCUMENTS_URL}/doc_123/files/page-citation-source").mock( + return_value=httpx.Response( + 200, + json={ + "document_id": "doc_123", + "namespace": "support-center", + "job_id": "job_123", + "job_result_id": "jres_123", + "variant": "normalized_pdf", + "file_name": "report.pdf", + "content_type": "application/pdf", + "url": "https://assets.example/report.pdf", + "expires_at": "2026-01-01T00:00:00Z", + }, + ) + ) + + source = sync_client.documents.get_page_citation_source("doc_123") + + assert route.called + assert source.document_id == "doc_123" + assert source.job_result_id == "jres_123" + assert source.content_type == "application/pdf" + assert source.url == "https://assets.example/report.pdf" + + @respx.mock + def test_get_page_citation_source_surfaces_not_found(self, sync_client: Any) -> None: + from knowhere._exceptions import NotFoundError + + respx.get(f"{DOCUMENTS_URL}/doc-404/files/page-citation-source").mock( + return_value=httpx.Response( + 404, + json={"error": {"code": "NOT_FOUND", "message": "not found"}}, + ) + ) + + with pytest.raises(NotFoundError): + sync_client.documents.get_page_citation_source("doc-404") + + @respx.mock + @pytest.mark.asyncio + async def test_async_get_page_citation_source( + self, + async_client: Any, + ) -> None: + route = respx.get(f"{DOCUMENTS_URL}/doc_123/files/page-citation-source").mock( + return_value=httpx.Response( + 200, + json={ + "document_id": "doc_123", + "file_name": "report.pdf", + "content_type": "application/pdf", + "url": "https://assets.example/report.pdf", + }, + ) + ) + + source = await async_client.documents.get_page_citation_source("doc_123") + + assert route.called + assert source.file_name == "report.pdf" + + def test_page_citation_asset_round_trips_descriptor(self) -> None: + from knowhere.types.page_citation import ( + PAGE_CITATION_ASSETS_METADATA_KEY, + PageCitationAsset, + ) + + payload = { + "page_num": 4, + "artifact_ref": "page_citation_assets/page-4.png", + "asset_url": "https://assets.example/page-4.png", + "content_type": "image/png", + "width": 1200, + "height": 1600, + "source": "knowhere-rendered-page-citation-source", + } + asset = PageCitationAsset.model_validate(payload) + assert asset.page_num == 4 + assert asset.source == "knowhere-rendered-page-citation-source" + chunk_metadata = {PAGE_CITATION_ASSETS_METADATA_KEY: [asset.model_dump()]} + parsed = [ + PageCitationAsset.model_validate(item) + for item in chunk_metadata[PAGE_CITATION_ASSETS_METADATA_KEY] + ] + assert parsed[0].artifact_ref == "page_citation_assets/page-4.png" + @respx.mock @pytest.mark.asyncio async def test_async_archive_document_returns_archived_state( diff --git a/tests/test_jobs.py b/tests/test_jobs.py index c9ace62..db37296 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -50,7 +50,7 @@ def test_create_with_url_source( assert job.source_type == "url" assert job.status == "pending" assert job.namespace == "support-center" - assert not hasattr(job, "document_id") + assert job.document_id is None @respx.mock def test_create_with_file_source( @@ -102,6 +102,8 @@ def test_create_sends_correct_body( assert body["data_id"] == "my_data_id" assert body["namespace"] == "support-center" assert body["document_id"] == "doc_123" + assert body["document_metadata"]["created_by_client"] == "python-sdk" + assert "client_version" in body["document_metadata"] @respx.mock def test_create_sends_llm_config( @@ -171,6 +173,67 @@ def test_create_omits_llm_config_when_none( assert route.called body: Dict[str, Any] = json.loads(route.calls[0].request.read()) assert "llm_config" not in body + assert body["document_metadata"]["created_by_client"] == "python-sdk" + + + @respx.mock + def test_create_attaches_default_document_metadata( + self, + sync_client: Any, + ) -> None: + """Job create always sends official-client telemetry metadata.""" + from knowhere._version import __version__ + + response_body: Dict[str, Any] = { + "job_id": "job_meta_defaults", + "status": "pending", + "source_type": "url", + "document_id": "doc_planned", + } + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=response_body)) + + job = sync_client.jobs.create( + source_type="url", + source_url="https://example.com/doc.pdf", + ) + + body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert body["document_metadata"] == { + "created_by_client": "python-sdk", + "client_version": __version__, + } + assert job.document_id == "doc_planned" + + @respx.mock + def test_create_lets_caller_document_metadata_win( + self, + sync_client: Any, + ) -> None: + """Caller metadata keys override official defaults.""" + from knowhere._version import __version__ + + response_body: Dict[str, Any] = { + "job_id": "job_meta_override", + "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", + document_metadata={ + "created_by_client": "cli", + "title": "Report.pdf", + }, + ) + + body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert body["document_metadata"] == { + "created_by_client": "cli", + "client_version": __version__, + "title": "Report.pdf", + } # --------------------------------------------------------------------------- diff --git a/tests/test_models.py b/tests/test_models.py index 73ba699..559903d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -57,7 +57,7 @@ def test_from_dict_with_document_scope(self) -> None: } job: Job = Job(**data) assert job.namespace == "support-center" - assert "document_id" not in job.model_dump() + assert job.document_id == "doc_123" def test_from_dict_with_upload(self) -> None: data: Dict[str, Any] = { diff --git a/tests/test_parse.py b/tests/test_parse.py index 3559797..6f9ac1d 100644 --- a/tests/test_parse.py +++ b/tests/test_parse.py @@ -151,6 +151,53 @@ def test_parse_url_forwards_llm_config( body: Dict[str, Any] = json.loads(create_route.calls[0].request.read()) assert body["llm_config"] == llm_config + @respx.mock + def test_parse_url_forwards_document_metadata( + self, + sync_client: Any, + sample_zip_bytes: bytes, + ) -> None: + """parse() threads document_metadata into jobs.create POST body.""" + import json + + from knowhere._version import __version__ + + job_id: str = "job_meta_parse" + result_url: str = "https://storage.example.com/result.zip" + 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", + document_metadata={"title": "Report.pdf"}, + poll_interval=0.01, + verify_checksum=False, + ) + + body: Dict[str, Any] = json.loads(create_route.calls[0].request.read()) + assert body["document_metadata"] == { + "created_by_client": "python-sdk", + "client_version": __version__, + "title": "Report.pdf", + } + # --------------------------------------------------------------------------- # parse(file=Path(...))