diff --git a/README.md b/README.md index 782cc02..c55c262 100644 --- a/README.md +++ b/README.md @@ -23,13 +23,20 @@ import knowhere client = knowhere.Knowhere(api_key="sk_...") -result = client.parse(url="https://example.com/report.pdf") +result = client.parse( + url="https://example.com/report.pdf", +) print(result.statistics.total_chunks) print(result.full_markdown[:200]) for chunk in result.text_chunks: print(chunk.content[:80]) + +for page in result.page_chunks: + print(page.content_source) # "summary" + print(page.content[:120]) # page-level summary + print(page.metadata.page_nums) # [4, 5, 6] ``` ## Retrieval and document lifecycle @@ -59,6 +66,7 @@ After the job is done and published, query the canonical document content: response = client.retrieval.query( namespace="support-center", query="How do I reset Bluetooth pairing?", + chunk_types=["page"], top_k=5, channels=["path", "term"], filter_mode="keep", @@ -72,9 +80,11 @@ print(response.stop_reason) print(response.failure_reason) for reference in response.referenced_chunks: - print(reference.chunk_id, reference.document_id, reference.asset_url) + print(reference.chunk_id, reference.chunk_type, reference.content_source) + print(reference.metadata, reference.asset_url) for result in response.results: + print(result.chunk_id, result.chunk_type, result.content_source) print(result.content) print(result.score) print(result.source.source_file_name, result.source.section_path) @@ -96,7 +106,7 @@ chunks = client.documents.list_chunks( document_id, page=1, page_size=50, - chunk_type="image", + chunk_type="page", include_asset_urls=True, ) print(chunks.pagination.total) @@ -107,7 +117,8 @@ if chunks.chunks: include_asset_urls=True, ) print(chunk.chunk.content) - print(chunk.chunk.asset_url) # Requested 7-day URL for image/table chunks. + print(chunk.chunk.metadata.get("page_nums")) # Page citations. + print(chunk.chunk.asset_url) # Requested 7-day URL when available. client.documents.archive(document_id) ``` @@ -164,8 +175,15 @@ result = client.parse(url="https://example.com/report.pdf") # Text chunks for chunk in result.text_chunks: - print(chunk.keywords) - print(chunk.summary) + print(chunk.metadata.keywords) + print(chunk.metadata.summary) + +# Page chunks (v2 page-memory results) +for chunk in result.page_chunks: + print(chunk.content_source) # "summary" + print(chunk.content[:120]) + print(chunk.metadata.page_nums) # citation pages + print(chunk.metadata.entities) # Image chunks (raw bytes loaded from ZIP) for chunk in result.image_chunks: diff --git a/docs/usage.md b/docs/usage.md index b311032..c8956db 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -93,6 +93,17 @@ for chunk in result.text_chunks: print(chunk.content[:80]) ``` +Page-memory chunks are returned by the v2 API used by this SDK: + +```python +result = client.parse(url="https://example.com/report.pdf") + +for page in result.page_chunks: + print(page.content_source) # "summary" + print(page.content[:120]) # page-level summary + print(page.metadata.page_nums) # [4, 5, 6] +``` + --- ## Parsing Documents @@ -207,6 +218,7 @@ result.statistics.total_chunks # 152 result.statistics.text_chunks # 120 result.statistics.image_chunks # 22 result.statistics.table_chunks # 10 +result.statistics.page_chunks # 48 result.statistics.total_pages # 48 # Full markdown — the entire document as markdown @@ -219,6 +231,7 @@ print(len(result.chunks)) # 152 result.text_chunks # List[TextChunk] result.image_chunks # List[ImageChunk] result.table_chunks # List[TableChunk] +result.page_chunks # List[PageChunk] # Lookup by ID chunk = result.getChunk("chunk_42") @@ -248,8 +261,8 @@ result.save("./output/report/") ## Chunk Types -Every chunk shares a base set of fields (`chunk_id`, `type`, `content`, `path`, -`metadata`). Worker metadata is kept in the `metadata` dict — it is **not** +Every chunk shares a base set of fields (`chunk_id`, `type`, `content_source`, +`content`, `path`, `metadata`). Worker metadata is kept in the `metadata` model — it is **not** flattened to top-level chunk properties. ### Base fields (all chunk types) @@ -257,10 +270,11 @@ flattened to top-level chunk properties. | Field | Type | Description | |-------|------|-------------| | `chunk_id` | `str` | Unique identifier | -| `type` | `str` | `"text"`, `"image"`, or `"table"` | +| `type` | `str` | `"text"`, `"image"`, `"table"`, or `"page"` | +| `content_source` | `str \| None` | What `content` represents. Page chunks usually use `"summary"`; text chunks default to `"content"`. | | `content` | `str` | Text content or placeholder | | `path` | `str \| None` | Document structure path | -| `metadata` | `dict` | Raw worker metadata (tokens, keywords, summary, length, page_nums, etc.) | +| `metadata` | `ChunkMetadata` | Raw worker metadata (tokens, keywords, summary, length, page_nums, entities, etc.) | ### TextChunk @@ -268,8 +282,8 @@ flattened to top-level chunk properties. for chunk in result.text_chunks: print(f"[{chunk.chunk_id}] {chunk.content[:60]}...") # Metadata is in chunk.metadata, not flattened: - keywords = chunk.metadata.get("keywords", []) - summary = chunk.metadata.get("summary") + keywords = chunk.metadata.keywords or [] + summary = chunk.metadata.summary if keywords: print(f" Keywords: {', '.join(keywords)}") if summary: @@ -306,6 +320,20 @@ for tbl in result.table_chunks: tbl.save("./output/tables/") # writes HTML file to disk ``` +### PageChunk + +Page chunks are returned by the v2 page-memory API. Their `content` is the +page-level summary, and `metadata.page_nums` gives the citation pages. + +```python +for page in result.page_chunks: + pages = page.metadata.page_nums or [] + print(page.content_source) # "summary" + print(page.content) + print(pages) + print(page.metadata.entities) +``` + --- ## Step-by-Step Control (Jobs API) @@ -412,6 +440,8 @@ The poller uses adaptive backoff: it starts at `poll_interval` and gradually inc result = client.jobs.load(job_result) # or pass a URL directly: result = client.jobs.load("https://storage.example.com/result.zip") +# or resolve a job id before loading: +result = client.jobs.load("job_123") ``` --- @@ -470,6 +500,7 @@ retryable `ConflictError` using the server error code `ABORTED`. response = client.retrieval.query( namespace="support-center", query="How do I pair a Bluetooth headset?", + chunk_types=["page"], top_k=5, ) @@ -486,12 +517,14 @@ print(response.evidence_text) # rendered evidence context, when returned print(response.stop_reason) # agentic termination reason, when returned print(response.failure_reason) # no-answer reason, when returned for ref in response.referenced_chunks: - print(ref.chunk_id, ref.document_id, ref.chunk_type) - print(ref.section_path, ref.file_path, ref.job_id, ref.asset_url) + print(ref.chunk_id, ref.document_id, ref.chunk_type, ref.content_source) + print(ref.section_path, ref.file_path, ref.job_id, ref.asset_url, ref.metadata) # Legacy results are always available for result in response.results: + print(result.chunk_id) print(result.content) + print(result.content_source) print(result.score) print(result.source.document_id) print(result.source.source_file_name) @@ -501,6 +534,8 @@ for result in response.results: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `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`. | Retrieval results expose `content`, not the older parse-result `text` field. Media results may include `asset_url` when the server can sign the referenced @@ -510,9 +545,12 @@ Each retrieval result uses one canonical source reference shape: ```python result.content +result.chunk_id # Optional[str] result.chunk_type +result.content_source result.score result.asset_url # Optional[str] +result.metadata # Optional[dict] result.source.document_id result.source.source_file_name result.source.section_path diff --git a/pyproject.toml b/pyproject.toml index bcc3c2a..e3d1abf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "knowhere-python-sdk" -version = "0.6.0" +version = "2.0.0" description = "Official Python SDK for the Knowhere document parsing API" readme = "README.md" license = "MIT" @@ -32,7 +32,7 @@ dependencies = [ [project.optional-dependencies] dev = [ - "pytest>=7.0.0", + "pytest>=7.0.0,<9.0.0", "pytest-asyncio>=0.23.0", "respx>=0.21.0", "ruff>=0.1.0", diff --git a/src/knowhere/__init__.py b/src/knowhere/__init__.py index 2b963d1..fa15df5 100644 --- a/src/knowhere/__init__.py +++ b/src/knowhere/__init__.py @@ -49,6 +49,7 @@ from knowhere.types.params import ParsingParams, WebhookConfig from knowhere.types.retrieval import ( RetrievalChannel, + RetrievalChunkType, RetrievalFilterMode, RetrievalReferencedChunk, RetrievalSectionExclusion, @@ -64,6 +65,7 @@ ImageChunk, ImageFileInfo, Manifest, + PageChunk, ParseResult, ProcessingCost, ProcessingMetadata, @@ -117,6 +119,7 @@ "DocumentListResponse", # Retrieval types "RetrievalChannel", + "RetrievalChunkType", "RetrievalFilterMode", "RetrievalReferencedChunk", "RetrievalSectionExclusion", @@ -126,6 +129,7 @@ # Result types "ParseResult", "Manifest", + "PageChunk", "Statistics", "Checksum", "FileIndex", diff --git a/src/knowhere/_base_client.py b/src/knowhere/_base_client.py index 815bf39..5994ff9 100644 --- a/src/knowhere/_base_client.py +++ b/src/knowhere/_base_client.py @@ -37,11 +37,13 @@ _logger = getLogger() # Error codes that are always safe to retry (matches server ALWAYS_RETRYABLE_ERROR_CODES) -_ALWAYS_RETRYABLE_ERROR_CODES: frozenset[str] = frozenset({ - "ABORTED", # 409 - Concurrency conflict - "UNAVAILABLE", # 503 - Service temporarily down - "DEADLINE_EXCEEDED", # 504 - Timeout -}) +_ALWAYS_RETRYABLE_ERROR_CODES: frozenset[str] = frozenset( + { + "ABORTED", # 409 - Concurrency conflict + "UNAVAILABLE", # 503 - Service temporarily down + "DEADLINE_EXCEEDED", # 504 - Timeout + } +) # RESOURCE_EXHAUSTED (429) is conditionally retryable: # - Rate limit: details.retry_after present → RETRY @@ -83,18 +85,12 @@ def __init__( f"or the {ENV_API_KEY} environment variable." ) self.api_key = resolved_key - self.base_url = ( - base_url - or os.environ.get(ENV_BASE_URL) - or DEFAULT_BASE_URL - ).rstrip("/") + 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 = ( upload_timeout if upload_timeout is not None else DEFAULT_UPLOAD_TIMEOUT ) - self.max_retries = ( - max_retries if max_retries is not None else DEFAULT_MAX_RETRIES - ) + 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]: @@ -112,13 +108,11 @@ def _buildRequestUrl(self, path: str) -> str: if path.startswith("http://") or path.startswith("https://"): return path clean_path: str = path.lstrip("/") - if not clean_path.startswith(API_VERSION): + if not clean_path.startswith("v2/"): clean_path = f"{API_VERSION}/{clean_path}" return f"{self.base_url}/{clean_path}" - def _parseErrorResponse( - self, response: httpx.Response - ) -> Optional[Dict[str, Any]]: + def _parseErrorResponse(self, response: httpx.Response) -> Optional[Dict[str, Any]]: """Try to parse a JSON error body; return ``None`` on failure.""" try: return response.json() # type: ignore[no-any-return] @@ -201,7 +195,7 @@ def _calculateRetryDelay( if retry_after is not None and retry_after > 0: return retry_after # Exponential backoff: 0.5 * 2^attempt, capped at 30s - base_delay: float = min(0.5 * (2 ** attempt), 30.0) + base_delay: float = min(0.5 * (2**attempt), 30.0) jitter: float = random.uniform(0, base_delay * 0.25) return base_delay + jitter @@ -304,21 +298,15 @@ def _request( continue raise APIConnectionError(str(exc)) from exc - _logger.debug( - "Response: %d %s", response.status_code, url - ) + _logger.debug("Response: %d %s", response.status_code, url) # Success if response.is_success: - api_response: APIResponse[T] = APIResponse( - response, cast_to - ) + api_response: APIResponse[T] = APIResponse(response, cast_to) return api_response.parse() # Error — decide whether to retry - error_body: Optional[Dict[str, Any]] = self._parseErrorResponse( - response - ) + error_body: Optional[Dict[str, Any]] = self._parseErrorResponse(response) error_code: Optional[str] = None error_details: Optional[Dict[str, Any]] = None if isinstance(error_body, dict): @@ -329,15 +317,10 @@ def _request( if isinstance(raw_details, dict): error_details = raw_details - if ( - attempt < self.max_retries - and self._shouldRetry( - response.status_code, error_code, error_details - ) + if attempt < self.max_retries and self._shouldRetry( + response.status_code, error_code, error_details ): - retry_after_val: Optional[float] = self._extractRetryAfter( - error_body, response - ) + retry_after_val: Optional[float] = self._extractRetryAfter(error_body, response) delay = self._calculateRetryDelay(attempt, retry_after_val) _logger.warning( "Retryable error %d on attempt %d/%d, retrying in %.1fs", @@ -441,7 +424,9 @@ async def _request( delay: float = self._calculateRetryDelay(attempt) _logger.warning( "Timeout on attempt %d/%d, retrying in %.1fs", - attempt + 1, self.max_retries + 1, delay, + attempt + 1, + self.max_retries + 1, + delay, ) await asyncio.sleep(delay) continue @@ -453,7 +438,9 @@ async def _request( delay = self._calculateRetryDelay(attempt) _logger.warning( "Connection error on attempt %d/%d, retrying in %.1fs", - attempt + 1, self.max_retries + 1, delay, + attempt + 1, + self.max_retries + 1, + delay, ) await asyncio.sleep(delay) continue @@ -476,19 +463,17 @@ async def _request( if isinstance(raw_details, dict): error_details = raw_details - if ( - attempt < self.max_retries - and self._shouldRetry( - response.status_code, error_code, error_details - ) + if attempt < self.max_retries and self._shouldRetry( + response.status_code, error_code, error_details ): - retry_after_val: Optional[float] = self._extractRetryAfter( - error_body, response - ) + retry_after_val: Optional[float] = self._extractRetryAfter(error_body, response) delay = self._calculateRetryDelay(attempt, retry_after_val) _logger.warning( "Retryable error %d on attempt %d/%d, retrying in %.1fs", - response.status_code, attempt + 1, self.max_retries + 1, delay, + response.status_code, + attempt + 1, + self.max_retries + 1, + delay, ) await asyncio.sleep(delay) continue diff --git a/src/knowhere/_client.py b/src/knowhere/_client.py index b45bdc1..dff40af 100644 --- a/src/knowhere/_client.py +++ b/src/knowhere/_client.py @@ -153,7 +153,10 @@ def parse( ) # Load and return parsed result - return self.jobs.load(job_result, verify_checksum=verify_checksum) + return self.jobs.load( + job_result, + verify_checksum=verify_checksum, + ) class AsyncKnowhere(AsyncAPIClient): @@ -273,5 +276,6 @@ async def parse( ) return await self.jobs.load( - job_result, verify_checksum=verify_checksum + job_result, + verify_checksum=verify_checksum, ) diff --git a/src/knowhere/_constants.py b/src/knowhere/_constants.py index 7b74dd3..3aba787 100644 --- a/src/knowhere/_constants.py +++ b/src/knowhere/_constants.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Literal + # Base URL for the Knowhere API DEFAULT_BASE_URL: str = "https://api.knowhereto.ai" @@ -26,7 +28,7 @@ POLL_BACKOFF_THRESHOLD: float = 60.0 # API version prefix -API_VERSION: str = "v1" +API_VERSION: Literal["v2"] = "v2" # Terminal job statuses that indicate polling should stop TERMINAL_STATUSES: frozenset[str] = frozenset({"done", "failed"}) diff --git a/src/knowhere/_version.py b/src/knowhere/_version.py index 42c2d87..e69b2f7 100644 --- a/src/knowhere/_version.py +++ b/src/knowhere/_version.py @@ -1 +1 @@ -__version__ = "0.6.0" # x-release-please-version +__version__ = "2.0.0" # x-release-please-version diff --git a/src/knowhere/lib/polling.py b/src/knowhere/lib/polling.py index e57455b..332c009 100644 --- a/src/knowhere/lib/polling.py +++ b/src/knowhere/lib/polling.py @@ -6,6 +6,7 @@ from typing import Optional from knowhere._constants import ( + API_VERSION, DEFAULT_POLL_INTERVAL, DEFAULT_POLL_TIMEOUT, MAX_POLL_INTERVAL, @@ -50,7 +51,7 @@ def syncPoll( poll_timeout: float = DEFAULT_POLL_TIMEOUT, on_progress: Optional[PollProgressCallback] = None, ) -> JobResult: - """Poll ``GET /v1/jobs/{job_id}`` until a terminal status is reached. + """Poll ``GET /v2/jobs/{job_id}`` until a terminal status is reached. Uses adaptive backoff: after ``POLL_BACKOFF_THRESHOLD`` seconds the interval grows by ``POLL_BACKOFF_MULTIPLIER``, capped at @@ -70,7 +71,7 @@ def syncPoll( job_result: JobResult = client._request( "GET", - f"v1/jobs/{job_id}", + f"{API_VERSION}/jobs/{job_id}", cast_to=JobResult, ) @@ -120,7 +121,7 @@ async def asyncPoll( job_result: JobResult = await client._request( "GET", - f"v1/jobs/{job_id}", + f"{API_VERSION}/jobs/{job_id}", cast_to=JobResult, ) diff --git a/src/knowhere/lib/result_parser.py b/src/knowhere/lib/result_parser.py index eac4579..e58b9b0 100644 --- a/src/knowhere/lib/result_parser.py +++ b/src/knowhere/lib/result_parser.py @@ -16,6 +16,7 @@ DocNav, ImageChunk, Manifest, + PageChunk, ParseResult, SlimChunk, TableChunk, @@ -39,9 +40,7 @@ def _safeZipPath(member_name: str, target_dir: str) -> str: """Validate that a ZIP member path does not escape *target_dir* (Zip Slip).""" abs_path: str = os.path.normpath(os.path.join(target_dir, member_name)) if not abs_path.startswith(os.path.normpath(target_dir)): - raise KnowhereError( - f"Zip Slip detected: '{member_name}' escapes target directory." - ) + raise KnowhereError(f"Zip Slip detected: '{member_name}' escapes target directory.") return abs_path @@ -81,6 +80,31 @@ def _extractFilePath(raw: Dict[str, Any]) -> Optional[str]: return fallback +def _extractContentSource( + raw: Dict[str, Any], + default: Optional[str] = None, +) -> Optional[str]: + """Return the chunk content source from snake_case or camelCase payloads.""" + raw_content_source: Any = raw.get("content_source", raw.get("contentSource")) + if isinstance(raw_content_source, str) and raw_content_source: + return raw_content_source + return default + + +def _extractPageContent(raw: Dict[str, Any]) -> str: + """Return page display content, preferring explicit content then summary.""" + raw_content: Any = raw.get("content") + if isinstance(raw_content, str) and raw_content: + return raw_content + + metadata: Any = raw.get("metadata", {}) + if isinstance(metadata, dict): + raw_summary: Any = metadata.get("summary") + if isinstance(raw_summary, str): + return raw_summary + return "" + + def _buildChunks( raw_chunks: List[Dict[str, Any]], zf: zipfile.ZipFile, @@ -90,15 +114,26 @@ def _buildChunks( for raw in raw_chunks: chunk_type: str = raw.get("type", "text") + chunk: Chunk - if chunk_type == "image": + if chunk_type == "page": + chunk = PageChunk( + chunk_id=raw.get("chunk_id", ""), + type="page", + content_source=_extractContentSource(raw, "summary"), + content=_extractPageContent(raw), + path=raw.get("path"), + metadata=raw.get("metadata", {}), + ) + elif chunk_type == "image": image_data: bytes = b"" file_path: Optional[str] = _extractFilePath(raw) if file_path: image_data = _readZipBytes(zf, file_path) or b"" - chunk: Chunk = ImageChunk( + chunk = ImageChunk( chunk_id=raw.get("chunk_id", ""), type="image", + content_source=_extractContentSource(raw), content=raw.get("content", ""), path=raw.get("path"), file_path=file_path, @@ -113,6 +148,7 @@ def _buildChunks( chunk = TableChunk( chunk_id=raw.get("chunk_id", ""), type="table", + content_source=_extractContentSource(raw), content=raw.get("content", ""), path=raw.get("path"), file_path=file_path, @@ -123,6 +159,7 @@ def _buildChunks( chunk = TextChunk( chunk_id=raw.get("chunk_id", ""), type="text", + content_source=_extractContentSource(raw, "content"), content=raw.get("content", ""), path=raw.get("path"), metadata=raw.get("metadata", {}), @@ -181,16 +218,12 @@ def parseResultZip( # -- DocNav (current worker output) -- doc_nav_text: Optional[str] = _readZipText(zf, "doc_nav.json") doc_nav: Optional[DocNav] = ( - DocNav.model_validate(json.loads(doc_nav_text)) - if doc_nav_text - else None + DocNav.model_validate(json.loads(doc_nav_text)) if doc_nav_text else None ) # -- Hierarchy (legacy — current worker no longer emits this) -- hierarchy_text: Optional[str] = _readZipText(zf, "hierarchy.json") - hierarchy: Optional[Any] = ( - json.loads(hierarchy_text) if hierarchy_text else None - ) + hierarchy: Optional[Any] = json.loads(hierarchy_text) if hierarchy_text else None # -- Optimized sidecar files -- chunks_slim_text: Optional[str] = _readZipText(zf, "chunks_slim.json") diff --git a/src/knowhere/resources/_base.py b/src/knowhere/resources/_base.py index f8a53fb..cdc30eb 100644 --- a/src/knowhere/resources/_base.py +++ b/src/knowhere/resources/_base.py @@ -4,6 +4,7 @@ from typing import Any, Dict, Optional, Type, TypeVar +from knowhere._constants import API_VERSION from knowhere._base_client import AsyncAPIClient, SyncAPIClient T = TypeVar("T") @@ -17,6 +18,9 @@ class SyncAPIResource: def __init__(self, client: SyncAPIClient) -> None: self._client = client + def _versionedPath(self, path: str) -> str: + return f"{API_VERSION}/{path.lstrip('/')}" + def _request( self, method: str, @@ -48,6 +52,9 @@ class AsyncAPIResource: def __init__(self, client: AsyncAPIClient) -> None: self._client = client + def _versionedPath(self, path: str) -> str: + return f"{API_VERSION}/{path.lstrip('/')}" + async def _request( self, method: str, diff --git a/src/knowhere/resources/documents.py b/src/knowhere/resources/documents.py index 7551dc7..e1cc183 100644 --- a/src/knowhere/resources/documents.py +++ b/src/knowhere/resources/documents.py @@ -15,7 +15,7 @@ class Documents(SyncAPIResource): - """Synchronous interface for ``/v1/documents`` endpoints.""" + """Synchronous interface for ``/v2/documents`` endpoints.""" def list( self, @@ -33,7 +33,7 @@ def list( return self._request( "GET", - "v1/documents", + self._versionedPath("documents"), params=params or None, cast_to=DocumentListResponse, ) @@ -42,7 +42,7 @@ def get(self, document_id: str) -> Document: """Get one canonical document by ID.""" return self._request( "GET", - f"v1/documents/{document_id}", + self._versionedPath(f"documents/{document_id}"), cast_to=Document, ) @@ -65,7 +65,7 @@ def list_chunks( return self._request( "GET", - f"v1/documents/{document_id}/chunks", + self._versionedPath(f"documents/{document_id}/chunks"), params=params or None, cast_to=DocumentChunkListResponse, ) @@ -84,7 +84,7 @@ def get_chunk( return self._request( "GET", - f"v1/documents/{document_id}/chunks/{document_chunk_id}", + self._versionedPath(f"documents/{document_id}/chunks/{document_chunk_id}"), params=params or None, cast_to=DocumentChunkResponse, ) @@ -93,13 +93,13 @@ def archive(self, document_id: str) -> Document: """Archive one canonical document by ID.""" return self._request( "POST", - f"v1/documents/{document_id}/archive", + self._versionedPath(f"documents/{document_id}/archive"), cast_to=Document, ) class AsyncDocuments(AsyncAPIResource): - """Asynchronous interface for ``/v1/documents`` endpoints.""" + """Asynchronous interface for ``/v2/documents`` endpoints.""" async def list( self, @@ -117,7 +117,7 @@ async def list( return await self._request( "GET", - "v1/documents", + self._versionedPath("documents"), params=params or None, cast_to=DocumentListResponse, ) @@ -126,7 +126,7 @@ async def get(self, document_id: str) -> Document: """Get one canonical document by ID.""" return await self._request( "GET", - f"v1/documents/{document_id}", + self._versionedPath(f"documents/{document_id}"), cast_to=Document, ) @@ -149,7 +149,7 @@ async def list_chunks( return await self._request( "GET", - f"v1/documents/{document_id}/chunks", + self._versionedPath(f"documents/{document_id}/chunks"), params=params or None, cast_to=DocumentChunkListResponse, ) @@ -168,7 +168,7 @@ async def get_chunk( return await self._request( "GET", - f"v1/documents/{document_id}/chunks/{document_chunk_id}", + self._versionedPath(f"documents/{document_id}/chunks/{document_chunk_id}"), params=params or None, cast_to=DocumentChunkResponse, ) @@ -177,7 +177,7 @@ async def archive(self, document_id: str) -> Document: """Archive one canonical document by ID.""" return await self._request( "POST", - f"v1/documents/{document_id}/archive", + self._versionedPath(f"documents/{document_id}/archive"), cast_to=Document, ) diff --git a/src/knowhere/resources/jobs.py b/src/knowhere/resources/jobs.py index f8e184b..faaf061 100644 --- a/src/knowhere/resources/jobs.py +++ b/src/knowhere/resources/jobs.py @@ -26,7 +26,7 @@ class Jobs(SyncAPIResource): - """Synchronous interface for the ``/v1/jobs`` endpoints.""" + """Synchronous interface for the ``/v2/jobs`` endpoints.""" def create( self, @@ -71,11 +71,20 @@ def create( if webhook is not None: body["webhook"] = dict(webhook) - return self._request("POST", "v1/jobs", body=body, cast_to=Job) + return self._request( + "POST", + self._versionedPath("jobs"), + body=body, + cast_to=Job, + ) def get(self, job_id: str) -> JobResult: """Retrieve the current status and result of a job.""" - return self._request("GET", f"v1/jobs/{job_id}", cast_to=JobResult) + return self._request( + "GET", + self._versionedPath(f"jobs/{job_id}"), + cast_to=JobResult, + ) def upload( self, @@ -148,9 +157,17 @@ def load( namespace: Optional[str] = job_result.namespace document_id: Optional[str] = job_result.document_id else: - result_url = job_result - namespace = None - document_id = None + if job_result.startswith("http://") or job_result.startswith("https://"): + result_url = job_result + namespace = None + document_id = None + else: + resolved_job_result = self.get(job_result) + if not resolved_job_result.result_url: + raise InvalidStateError("JobResult does not have a result_url.") + result_url = resolved_job_result.result_url + namespace = resolved_job_result.namespace + document_id = resolved_job_result.document_id response: httpx.Response = self._client._client.get( result_url, timeout=self._client.upload_timeout @@ -165,7 +182,7 @@ def load( class AsyncJobs(AsyncAPIResource): - """Asynchronous interface for the ``/v1/jobs`` endpoints.""" + """Asynchronous interface for the ``/v2/jobs`` endpoints.""" async def create( self, @@ -196,12 +213,19 @@ async def create( if webhook is not None: body["webhook"] = dict(webhook) - return await self._request("POST", "v1/jobs", body=body, cast_to=Job) + return await self._request( + "POST", + self._versionedPath("jobs"), + body=body, + cast_to=Job, + ) async def get(self, job_id: str) -> JobResult: """Retrieve the current status and result of a job (async).""" return await self._request( - "GET", f"v1/jobs/{job_id}", cast_to=JobResult + "GET", + self._versionedPath(f"jobs/{job_id}"), + cast_to=JobResult, ) async def upload( @@ -261,9 +285,17 @@ async def load( namespace: Optional[str] = job_result.namespace document_id: Optional[str] = job_result.document_id else: - result_url = job_result - namespace = None - document_id = None + if job_result.startswith("http://") or job_result.startswith("https://"): + result_url = job_result + namespace = None + document_id = None + else: + resolved_job_result = await self.get(job_result) + if not resolved_job_result.result_url: + raise InvalidStateError("JobResult does not have a result_url.") + result_url = resolved_job_result.result_url + namespace = resolved_job_result.namespace + document_id = resolved_job_result.document_id response: httpx.Response = await self._client._client.get( result_url, timeout=self._client.upload_timeout diff --git a/src/knowhere/resources/retrieval.py b/src/knowhere/resources/retrieval.py index 9100be8..ad3e940 100644 --- a/src/knowhere/resources/retrieval.py +++ b/src/knowhere/resources/retrieval.py @@ -6,6 +6,7 @@ from knowhere.resources._base import AsyncAPIResource, SyncAPIResource from knowhere.types.retrieval import ( + RetrievalChunkType, RetrievalChannel, RetrievalFilterMode, RetrievalQueryResponse, @@ -14,7 +15,7 @@ class Retrieval(SyncAPIResource): - """Synchronous interface for ``/v1/retrieval`` endpoints.""" + """Synchronous interface for ``/v2/retrieval`` endpoints.""" def query( self, @@ -24,6 +25,7 @@ def query( top_k: Optional[int] = None, use_agentic: Optional[bool] = None, data_type: Optional[int] = None, + chunk_types: Optional[list[RetrievalChunkType]] = None, signal_paths: Optional[list[str]] = None, filter_mode: Optional[RetrievalFilterMode] = None, channels: Optional[list[RetrievalChannel]] = None, @@ -44,6 +46,8 @@ def query( body["use_agentic"] = use_agentic if data_type is not None: body["data_type"] = data_type + if chunk_types is not None: + body["chunk_types"] = chunk_types if signal_paths is not None: body["signal_paths"] = signal_paths if filter_mode is not None: @@ -65,14 +69,14 @@ def query( return self._request( "POST", - "v1/retrieval/query", + self._versionedPath("retrieval/query"), body=body, cast_to=RetrievalQueryResponse, ) class AsyncRetrieval(AsyncAPIResource): - """Asynchronous interface for ``/v1/retrieval`` endpoints.""" + """Asynchronous interface for ``/v2/retrieval`` endpoints.""" async def query( self, @@ -82,6 +86,7 @@ async def query( top_k: Optional[int] = None, use_agentic: Optional[bool] = None, data_type: Optional[int] = None, + chunk_types: Optional[list[RetrievalChunkType]] = None, signal_paths: Optional[list[str]] = None, filter_mode: Optional[RetrievalFilterMode] = None, channels: Optional[list[RetrievalChannel]] = None, @@ -102,6 +107,8 @@ async def query( body["use_agentic"] = use_agentic if data_type is not None: body["data_type"] = data_type + if chunk_types is not None: + body["chunk_types"] = chunk_types if signal_paths is not None: body["signal_paths"] = signal_paths if filter_mode is not None: @@ -123,7 +130,7 @@ async def query( return await self._request( "POST", - "v1/retrieval/query", + self._versionedPath("retrieval/query"), body=body, cast_to=RetrievalQueryResponse, ) diff --git a/src/knowhere/types/__init__.py b/src/knowhere/types/__init__.py index 14274c2..6588f6f 100644 --- a/src/knowhere/types/__init__.py +++ b/src/knowhere/types/__init__.py @@ -16,6 +16,7 @@ from knowhere.types.params import ParsingParams, WebhookConfig from knowhere.types.retrieval import ( RetrievalChannel, + RetrievalChunkType, RetrievalFilterMode, RetrievalReferencedChunk, RetrievalSectionExclusion, @@ -31,6 +32,7 @@ ImageChunk, ImageFileInfo, Manifest, + PageChunk, ParseResult, ProcessingCost, ProcessingMetadata, @@ -58,6 +60,7 @@ "DocumentListResponse", # retrieval "RetrievalChannel", + "RetrievalChunkType", "RetrievalFilterMode", "RetrievalReferencedChunk", "RetrievalSectionExclusion", @@ -75,6 +78,7 @@ "ImageChunk", "ImageFileInfo", "Manifest", + "PageChunk", "ParseResult", "ProcessingCost", "ProcessingMetadata", diff --git a/src/knowhere/types/document.py b/src/knowhere/types/document.py index e6e813c..577f03e 100644 --- a/src/knowhere/types/document.py +++ b/src/knowhere/types/document.py @@ -9,7 +9,7 @@ class Document(BaseModel): - """Canonical document state returned by ``/v1/documents`` endpoints.""" + """Canonical document state returned by ``/v2/documents`` endpoints.""" document_id: str namespace: str @@ -31,7 +31,7 @@ class DocumentListPagination(BaseModel): class DocumentListResponse(BaseModel): - """Response from ``GET /v1/documents``.""" + """Response from ``GET /v2/documents``.""" namespace: str documents: list[Document] @@ -55,7 +55,7 @@ def populate_legacy_pagination(cls, value: Any) -> Any: return payload -DocumentChunkType = Literal["text", "image", "table"] +DocumentChunkType = Literal["text", "image", "table", "page"] class DocumentChunkPagination(BaseModel): @@ -73,6 +73,7 @@ class DocumentChunk(BaseModel): id: str chunk_id: str chunk_type: DocumentChunkType + content_source: Optional[str] = None content: Optional[str] = None section_id: Optional[str] = None section_path: Optional[str] = None @@ -85,7 +86,7 @@ class DocumentChunk(BaseModel): class DocumentChunkListResponse(BaseModel): - """Response from ``GET /v1/documents/{document_id}/chunks``.""" + """Response from ``GET /v2/documents/{document_id}/chunks``.""" document_id: str namespace: str @@ -96,7 +97,7 @@ class DocumentChunkListResponse(BaseModel): class DocumentChunkResponse(BaseModel): - """Response from ``GET /v1/documents/{document_id}/chunks/{chunk_id}``.""" + """Response from ``GET /v2/documents/{document_id}/chunks/{chunk_id}``.""" document_id: str namespace: str diff --git a/src/knowhere/types/job.py b/src/knowhere/types/job.py index ea86556..537a562 100644 --- a/src/knowhere/types/job.py +++ b/src/knowhere/types/job.py @@ -35,7 +35,7 @@ def fraction(self) -> float: class Job(BaseModel): - """Response from ``POST /v1/jobs`` — represents a newly created job.""" + """Response from ``POST /v2/jobs`` — represents a newly created job.""" job_id: str status: str @@ -49,7 +49,7 @@ class Job(BaseModel): class JobResult(BaseModel): - """Response from ``GET /v1/jobs/{job_id}`` — full job status and result.""" + """Response from ``GET /v2/jobs/{job_id}`` — full job status and result.""" job_id: str status: str diff --git a/src/knowhere/types/params.py b/src/knowhere/types/params.py index 9882fca..8359066 100644 --- a/src/knowhere/types/params.py +++ b/src/knowhere/types/params.py @@ -2,7 +2,6 @@ from __future__ import annotations - from typing_extensions import TypedDict diff --git a/src/knowhere/types/result.py b/src/knowhere/types/result.py index 0cce2da..77c8adf 100644 --- a/src/knowhere/types/result.py +++ b/src/knowhere/types/result.py @@ -55,6 +55,7 @@ class Statistics(BaseModel): text_chunks: Optional[int] = 0 image_chunks: Optional[int] = 0 table_chunks: Optional[int] = 0 + page_chunks: Optional[int] = 0 total_pages: Optional[int] = 0 @@ -193,6 +194,7 @@ class ChunkMetadata(BaseModel): length: Optional[int] = None page_nums: Optional[List[int]] = None + entities: Optional[List[Dict[str, Any]]] = None tokens: Optional[List[str]] = None keywords: Optional[List[str]] = None summary: Optional[str] = None @@ -208,6 +210,7 @@ class BaseChunk(BaseModel): chunk_id: str type: str + content_source: Optional[str] = None content: str = "" path: Optional[str] = None metadata: ChunkMetadata = Field(default_factory=ChunkMetadata) @@ -245,13 +248,9 @@ def save(self, directory: Union[str, Path]) -> Path: dir_path: Path = Path(directory) dir_path.mkdir(parents=True, exist_ok=True) - raw_name: str = os.path.basename( - self.file_path or f"{self.chunk_id}.bin" - ) + raw_name: str = os.path.basename(self.file_path or f"{self.chunk_id}.bin") safe_name: str = _sanitizeFilename(raw_name) - out_path: Path = _ensurePathWithinDirectory( - dir_path, dir_path / safe_name - ) + out_path: Path = _ensurePathWithinDirectory(dir_path, dir_path / safe_name) out_path.write_bytes(self.data) return out_path @@ -268,19 +267,21 @@ def save(self, directory: Union[str, Path]) -> Path: dir_path: Path = Path(directory) dir_path.mkdir(parents=True, exist_ok=True) - raw_name: str = os.path.basename( - self.file_path or f"{self.chunk_id}.html" - ) + raw_name: str = os.path.basename(self.file_path or f"{self.chunk_id}.html") safe_name: str = _sanitizeFilename(raw_name) - out_path: Path = _ensurePathWithinDirectory( - dir_path, dir_path / safe_name - ) + out_path: Path = _ensurePathWithinDirectory(dir_path, dir_path / safe_name) out_path.write_text(self.html, encoding="utf-8") return out_path +class PageChunk(BaseChunk): + """A page chunk. Its content usually contains a page-level summary.""" + + type: str = "page" + + # Union of all chunk types -Chunk = Union[TextChunk, ImageChunk, TableChunk] +Chunk = Union[TextChunk, ImageChunk, TableChunk, PageChunk] class SlimChunk(BaseModel): @@ -371,6 +372,11 @@ def table_chunks(self) -> List[TableChunk]: """Return only table chunks.""" return [c for c in self.chunks if isinstance(c, TableChunk)] + @property + def page_chunks(self) -> List[PageChunk]: + """Return only page chunks.""" + return [c for c in self.chunks if isinstance(c, PageChunk)] + @property def job_id(self) -> Optional[str]: """Shortcut to ``manifest.job_id``.""" diff --git a/src/knowhere/types/retrieval.py b/src/knowhere/types/retrieval.py index 640c63f..5e8fbb6 100644 --- a/src/knowhere/types/retrieval.py +++ b/src/knowhere/types/retrieval.py @@ -8,6 +8,7 @@ RetrievalChannel = Literal["path", "content", "term"] +RetrievalChunkType = Literal["text", "image", "table", "page"] RetrievalFilterMode = Literal["delete", "keep"] @@ -27,12 +28,17 @@ class RetrievalSource(BaseModel): class RetrievalResult(BaseModel): - """Canonical chunk result returned by ``POST /v1/retrieval/query``.""" + """Canonical chunk result returned by ``POST /v2/retrieval/query``.""" + chunk_id: Optional[str] = None chunk_type: str + content_source: Optional[str] = None content: str score: Optional[float] = None asset_url: Optional[str] = None + source_chunk_path: Optional[str] = None + file_path: Optional[str] = None + metadata: Optional[dict[str, Any]] = None source: RetrievalSource @@ -42,14 +48,17 @@ class RetrievalReferencedChunk(BaseModel): chunk_id: str document_id: str chunk_type: str + content_source: Optional[str] = None section_path: str + source_chunk_path: Optional[str] = None file_path: Optional[str] = None job_id: Optional[str] = None asset_url: Optional[str] = None + metadata: Optional[dict[str, Any]] = None class RetrievalQueryResponse(BaseModel): - """Response from ``POST /v1/retrieval/query``. + """Response from ``POST /v2/retrieval/query``. Three PRIMARY output fields for downstream agent consumption: diff --git a/tests/conftest.py b/tests/conftest.py index c742f61..2152b8c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,7 +62,7 @@ async def async_client(api_key: str, base_url: str) -> Any: @pytest.fixture() def mock_job_response() -> Dict[str, Any]: - """Return a dict matching POST /v1/jobs response for a file upload job. + """Return a dict matching POST /v2/jobs response for a file upload job. The APIResponse.parse() calls model_validate() on the raw JSON, so the response body must match the Job model directly. @@ -82,7 +82,7 @@ def mock_job_response() -> Dict[str, Any]: @pytest.fixture() def mock_job_result_response() -> Dict[str, Any]: - """Return a dict matching GET /v1/jobs/{id} for a completed job.""" + """Return a dict matching GET /v2/jobs/{id} for a completed job.""" return { "job_id": "job_test123", "status": "done", diff --git a/tests/test_documents.py b/tests/test_documents.py index eb6c7c8..94f71a8 100644 --- a/tests/test_documents.py +++ b/tests/test_documents.py @@ -11,7 +11,7 @@ from tests.conftest import BASE_URL -DOCUMENTS_URL: str = f"{BASE_URL}/v1/documents" +DOCUMENTS_URL: str = f"{BASE_URL}/v2/documents" def _make_document(status: str = "active") -> Dict[str, Any]: @@ -32,6 +32,7 @@ def _make_document_chunk(chunk_type: str = "text") -> Dict[str, Any]: "id": "dchk_123", "chunk_id": "parser-chunk-1", "chunk_type": chunk_type, + "content_source": "summary" if chunk_type == "page" else "content", "content": "Chunk content", "section_id": "sec_123", "section_path": "Chapter 1", @@ -175,6 +176,41 @@ def test_list_chunks_sends_optional_query_params(self, sync_client: Any) -> None assert response.chunks[0].id == "dchk_123" assert response.pagination.total_pages == 2 + @respx.mock + def test_list_chunks_supports_page_chunks( + self, + sync_client: Any, + ) -> None: + route = respx.get(f"{DOCUMENTS_URL}/doc_123/chunks").mock( + return_value=httpx.Response( + 200, + json={ + "document_id": "doc_123", + "namespace": "support-center", + "job_result_id": "result_123", + "job_id": "job_123", + "chunks": [_make_document_chunk(chunk_type="page")], + "pagination": { + "page": 1, + "page_size": 50, + "total": 1, + "total_pages": 1, + }, + }, + ) + ) + + response = sync_client.documents.list_chunks( + "doc_123", + chunk_type="page", + ) + + assert route.called + assert route.calls[0].request.url.params["chunk_type"] == "page" + assert response.chunks[0].chunk_type == "page" + assert response.chunks[0].content_source == "summary" + assert response.chunks[0].metadata["page_nums"] == [1] + @respx.mock def test_list_chunks_omits_default_query_params(self, sync_client: Any) -> None: route = respx.get(f"{DOCUMENTS_URL}/doc_123/chunks").mock( @@ -204,9 +240,7 @@ def test_list_chunks_omits_default_query_params(self, sync_client: Any) -> None: assert response.pagination.total == 0 @respx.mock - def test_list_chunks_accepts_explicit_asset_url_control( - self, sync_client: Any - ) -> None: + def test_list_chunks_accepts_explicit_asset_url_control(self, sync_client: Any) -> None: route = respx.get(f"{DOCUMENTS_URL}/doc_123/chunks").mock( return_value=httpx.Response( 200, diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py index 2c5ccc7..5a9cdd2 100644 --- a/tests/test_exceptions.py +++ b/tests/test_exceptions.py @@ -43,7 +43,7 @@ def _make_response( return httpx.Response( status_code=status_code, headers=headers or {}, - request=httpx.Request("GET", "https://api.test.knowhereto.ai/v1/jobs/j1"), + request=httpx.Request("GET", "https://api.test.knowhereto.ai/v2/jobs/j1"), ) diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 85669f1..6d09564 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from typing import Any, Dict import httpx @@ -14,8 +15,7 @@ # Helpers # --------------------------------------------------------------------------- -JOBS_URL: str = f"{BASE_URL}/v1/jobs" - +JOBS_URL: str = f"{BASE_URL}/v2/jobs" # --------------------------------------------------------------------------- # jobs.create() @@ -30,7 +30,7 @@ def test_create_with_url_source( self, sync_client: Any, ) -> None: - """POST /v1/jobs with source_type=url sends correct payload.""" + """POST /v2/jobs with source_type=url sends correct payload.""" response_body: Dict[str, Any] = { "job_id": "job_test123", "status": "pending", @@ -38,9 +38,7 @@ def test_create_with_url_source( "namespace": "support-center", } - route = respx.post(JOBS_URL).mock( - return_value=httpx.Response(200, json=response_body) - ) + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=response_body)) job = sync_client.jobs.create( source_type="url", @@ -60,10 +58,8 @@ def test_create_with_file_source( sync_client: Any, mock_job_response: Dict[str, Any], ) -> None: - """POST /v1/jobs with source_type=file returns upload_url.""" - route = respx.post(JOBS_URL).mock( - return_value=httpx.Response(200, json=mock_job_response) - ) + """POST /v2/jobs with source_type=file returns upload_url.""" + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=mock_job_response)) job = sync_client.jobs.create( source_type="file", @@ -88,9 +84,7 @@ def test_create_sends_correct_body( "namespace": "support-center", } - route = respx.post(JOBS_URL).mock( - return_value=httpx.Response(200, json=response_body) - ) + route = respx.post(JOBS_URL).mock(return_value=httpx.Response(200, json=response_body)) sync_client.jobs.create( source_type="url", @@ -101,8 +95,7 @@ def test_create_sends_correct_body( ) assert route.called - request_body: Dict[str, Any] = route.calls[0].request.read() - import json + request_body: bytes = route.calls[0].request.read() body: Dict[str, Any] = json.loads(request_body) assert body["source_type"] == "url" assert body["source_url"] == "https://example.com/doc.pdf" @@ -125,7 +118,7 @@ def test_get_returns_job_result( sync_client: Any, mock_job_result_response: Dict[str, Any], ) -> None: - """GET /v1/jobs/{job_id} returns a JobResult.""" + """GET /v2/jobs/{job_id} returns a JobResult.""" job_id: str = "job_test123" route = respx.get(f"{JOBS_URL}/{job_id}").mock( return_value=httpx.Response(200, json=mock_job_result_response) @@ -148,16 +141,12 @@ class TestJobsUpload: """Verify jobs.upload() sends PUT to the presigned URL.""" @respx.mock - def test_upload_sends_put_with_job_object( - self, sync_client: Any - ) -> None: + def test_upload_sends_put_with_job_object(self, sync_client: Any) -> None: """Upload sends PUT with file content using a Job object.""" from knowhere.types.job import Job upload_url: str = "https://storage.example.com/upload?token=abc" - route = respx.put(upload_url).mock( - return_value=httpx.Response(200) - ) + route = respx.put(upload_url).mock(return_value=httpx.Response(200)) job: Job = Job( job_id="job_upload", @@ -172,14 +161,10 @@ def test_upload_sends_put_with_job_object( assert route.called @respx.mock - def test_upload_sends_put_with_url_string( - self, sync_client: Any - ) -> None: + def test_upload_sends_put_with_url_string(self, sync_client: Any) -> None: """Upload sends PUT when given a URL string directly.""" upload_url: str = "https://storage.example.com/upload?token=def" - route = respx.put(upload_url).mock( - return_value=httpx.Response(200) - ) + route = respx.put(upload_url).mock(return_value=httpx.Response(200)) sync_client.jobs.upload(upload_url, b"fake pdf content") @@ -220,9 +205,7 @@ def test_wait_polls_until_done( ] ) - result = sync_client.jobs.wait( - job_id, poll_interval=0.01, poll_timeout=5.0 - ) + result = sync_client.jobs.wait(job_id, poll_interval=0.01, poll_timeout=5.0) assert result.status == "done" assert route.call_count == 2 @@ -252,9 +235,7 @@ def test_load_downloads_and_parses( ) ) - parse_result = sync_client.jobs.load( - result_url, verify_checksum=False - ) + parse_result = sync_client.jobs.load(result_url, verify_checksum=False) assert route.called assert parse_result.manifest is not None @@ -287,11 +268,48 @@ def test_load_with_job_result_object( result_url=result_url, ) - parse_result = sync_client.jobs.load( - job_result, verify_checksum=False - ) + parse_result = sync_client.jobs.load(job_result, verify_checksum=False) assert route.called assert parse_result.manifest is not None assert parse_result.namespace == "support-center" assert parse_result.document_id == "doc_123" + + @respx.mock + def test_load_resolves_job_id( + self, + sync_client: Any, + sample_zip_bytes: bytes, + ) -> None: + """Passing a job id resolves the result URL before loading the ZIP.""" + job_id: str = "job_v2" + result_url: str = "https://storage.example.com/result.zip" + status_route = respx.get(f"{JOBS_URL}/{job_id}").mock( + return_value=httpx.Response( + 200, + json={ + "job_id": job_id, + "status": "done", + "source_type": "url", + "namespace": "support-center", + "document_id": "doc_123", + "result_url": result_url, + }, + ) + ) + download_route = respx.get(result_url).mock( + return_value=httpx.Response( + 200, + content=sample_zip_bytes, + headers={"Content-Type": "application/zip"}, + ) + ) + + parse_result = sync_client.jobs.load( + job_id, + verify_checksum=False, + ) + + assert status_route.called + assert download_route.called + assert parse_result.namespace == "support-center" diff --git a/tests/test_models.py b/tests/test_models.py index 4314cfa..73ba699 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -9,10 +9,12 @@ BaseChunk, Checksum, Chunk, + ChunkMetadata, FileIndex, ImageChunk, ImageFileInfo, Manifest, + PageChunk, ParseResult, ProcessingCost, ProcessingMetadata, @@ -254,6 +256,7 @@ def test_statistics_accessible(self) -> None: text_chunks=3, image_chunks=1, table_chunks=1, + page_chunks=0, total_pages=2, ) ) @@ -266,12 +269,8 @@ def test_files_index(self) -> None: files=FileIndex( chunks="chunks.json", markdown="full.md", - images=[ - ImageFileInfo(id="IMG_1", file_path="images/IMG_1.jpg") - ], - tables=[ - TableFileInfo(id="TBL_1", file_path="tables/TBL_1.csv") - ], + images=[ImageFileInfo(id="IMG_1", file_path="images/IMG_1.jpg")], + tables=[TableFileInfo(id="TBL_1", file_path="tables/TBL_1.csv")], ) ) assert manifest.files is not None @@ -334,6 +333,7 @@ def test_defaults_to_zero(self) -> None: assert stats.text_chunks == 0 assert stats.image_chunks == 0 assert stats.table_chunks == 0 + assert stats.page_chunks == 0 assert stats.total_pages == 0 @@ -421,7 +421,7 @@ def test_metadata_accessible(self) -> None: chunk: BaseChunk = BaseChunk( chunk_id="chunk_3", type="text", - metadata={"tokens": ["a", "b"], "length": 10}, + metadata=ChunkMetadata(tokens=["a", "b"], length=10), ) assert chunk.metadata.tokens == ["a", "b"] assert chunk.metadata.length == 10 @@ -484,9 +484,7 @@ def test_defaults(self) -> None: assert chunk.data == b"" def test_format_property_from_file_path(self) -> None: - chunk: ImageChunk = ImageChunk( - chunk_id="IMG_3", file_path="images/IMG_3.png" - ) + chunk: ImageChunk = ImageChunk(chunk_id="IMG_3", file_path="images/IMG_3.png") assert chunk.format == "png" def test_format_property_none_without_file_path(self) -> None: @@ -498,9 +496,7 @@ def test_is_instance_of_base_chunk(self) -> None: assert isinstance(chunk, BaseChunk) def test_data_excluded_from_serialization(self) -> None: - chunk: ImageChunk = ImageChunk( - chunk_id="IMG_6", data=b"secret bytes" - ) + chunk: ImageChunk = ImageChunk(chunk_id="IMG_6", data=b"secret bytes") dumped: Dict[str, Any] = chunk.model_dump() assert "data" not in dumped @@ -535,13 +531,41 @@ def test_is_instance_of_base_chunk(self) -> None: assert isinstance(chunk, BaseChunk) def test_html_excluded_from_serialization(self) -> None: - chunk: TableChunk = TableChunk( - chunk_id="TBL_4", html="
" - ) + chunk: TableChunk = TableChunk(chunk_id="TBL_4", html="
") dumped: Dict[str, Any] = chunk.model_dump() assert "html" not in dumped +# --------------------------------------------------------------------------- +# PageChunk model +# --------------------------------------------------------------------------- + + +class TestPageChunkModel: + """Verify PageChunk fields and defaults.""" + + def test_from_dict(self) -> None: + chunk: PageChunk = PageChunk( + chunk_id="page_1", + content_source="summary", + content="Overview of pages 4-6", + metadata=ChunkMetadata( + page_nums=[4, 5, 6], + entities=[{"text": "refund", "label": "topic"}], + ), + ) + + assert chunk.chunk_id == "page_1" + assert chunk.type == "page" + assert chunk.content_source == "summary" + assert chunk.metadata.page_nums == [4, 5, 6] + assert chunk.metadata.entities == [{"text": "refund", "label": "topic"}] + + def test_is_instance_of_base_chunk(self) -> None: + chunk: PageChunk = PageChunk(chunk_id="page_2") + assert isinstance(chunk, BaseChunk) + + # --------------------------------------------------------------------------- # ParseResult # --------------------------------------------------------------------------- @@ -560,6 +584,7 @@ def _build_parse_result( text_chunks=1, image_chunks=1, table_chunks=1, + page_chunks=0, total_pages=2, ), ) @@ -658,8 +683,25 @@ def test_empty_chunks_list(self) -> None: assert len(result.text_chunks) == 0 assert len(result.image_chunks) == 0 assert len(result.table_chunks) == 0 + assert len(result.page_chunks) == 0 assert result.getChunk("anything") is None + def test_page_chunks_filters_correctly(self) -> None: + result: ParseResult = _build_parse_result( + chunks=[ + PageChunk( + chunk_id="page_1", + content_source="summary", + content="Page summary", + metadata=ChunkMetadata(page_nums=[1, 2]), + ) + ] + ) + + assert len(result.page_chunks) == 1 + assert result.page_chunks[0].chunk_id == "page_1" + assert result.page_chunks[0].metadata.page_nums == [1, 2] + def test_full_markdown_accessible(self) -> None: result: ParseResult = _build_parse_result() assert result.full_markdown == "# Test\n\nHello world" diff --git a/tests/test_parse.py b/tests/test_parse.py index c2f6c84..f12d6b2 100644 --- a/tests/test_parse.py +++ b/tests/test_parse.py @@ -15,7 +15,7 @@ from tests.conftest import BASE_URL -JOBS_URL: str = f"{BASE_URL}/v1/jobs" +JOBS_URL: str = f"{BASE_URL}/v2/jobs" def _make_create_response( @@ -23,7 +23,7 @@ def _make_create_response( source_type: str, upload_url: str | None = None, ) -> Dict[str, Any]: - """Build a mock POST /v1/jobs response (raw model data).""" + """Build a mock POST /v2/jobs response (raw model data).""" return { "job_id": job_id, "status": "waiting-file" if source_type == "file" else "pending", @@ -37,7 +37,7 @@ def _make_create_response( def _make_done_response(job_id: str, result_url: str) -> Dict[str, Any]: - """Build a mock GET /v1/jobs/{id} response for a completed job.""" + """Build a mock GET /v2/jobs/{id} response for a completed job.""" return { "job_id": job_id, "status": "done", diff --git a/tests/test_polling.py b/tests/test_polling.py index 647f513..cf63c24 100644 --- a/tests/test_polling.py +++ b/tests/test_polling.py @@ -13,7 +13,7 @@ from tests.conftest import BASE_URL -JOBS_URL: str = f"{BASE_URL}/v1/jobs" +JOBS_URL: str = f"{BASE_URL}/v2/jobs" def _make_status_response( @@ -21,7 +21,7 @@ def _make_status_response( status: str, error: Dict[str, Any] | None = None, ) -> Dict[str, Any]: - """Build a mock GET /v1/jobs/{id} response with the given status.""" + """Build a mock GET /v2/jobs/{id} response with the given status.""" data: Dict[str, Any] = { "job_id": job_id, "status": status, diff --git a/tests/test_result_parser.py b/tests/test_result_parser.py index 3ff3f24..8494036 100644 --- a/tests/test_result_parser.py +++ b/tests/test_result_parser.py @@ -17,6 +17,7 @@ DocNav, ImageChunk, Manifest, + PageChunk, ParseResult, Statistics, TableChunk, @@ -127,6 +128,7 @@ def _make_manifest(checksum_value: str = "") -> Dict[str, Any]: "text_chunks": 1, "image_chunks": 1, "table_chunks": 0, + "page_chunks": 0, "total_pages": 1, }, "files": { @@ -174,6 +176,7 @@ def _make_optimized_manifest() -> Dict[str, Any]: "text_chunks": 1, "image_chunks": 1, "table_chunks": 1, + "page_chunks": 0, "total_pages": None, }, } @@ -316,9 +319,9 @@ def test_exposes_optimized_payload_metadata_and_sidecar_assets(self) -> None: } ).encode("utf-8"), "kb.csv": b"chunk_id,type\ntext_chunk_optimized,text\n", - "hierarchy.json": json.dumps( - {"Default_Root": {"optimized.pdf": {}}} - ).encode("utf-8"), + "hierarchy.json": json.dumps({"Default_Root": {"optimized.pdf": {}}}).encode( + "utf-8" + ), "toc_hierarchies.json": json.dumps( [{"toc_range": [1, 3], "scan_range": [1, 10]}] ).encode("utf-8"), @@ -367,9 +370,9 @@ def test_save_preserves_optimized_sidecar_files(self, tmp_path: Path) -> None: } ).encode("utf-8"), "kb.csv": b"chunk_id,type\ntext_chunk_optimized,text\n", - "hierarchy.json": json.dumps( - {"Default_Root": {"optimized.pdf": {}}} - ).encode("utf-8"), + "hierarchy.json": json.dumps({"Default_Root": {"optimized.pdf": {}}}).encode( + "utf-8" + ), "toc_hierarchies.json": json.dumps( [{"toc_range": [1, 3], "scan_range": [1, 10]}] ).encode("utf-8"), @@ -422,6 +425,7 @@ def _make_current_contract_manifest() -> Dict[str, Any]: "text_chunks": 1, "image_chunks": 1, "table_chunks": 0, + "page_chunks": 0, "total_pages": None, }, } @@ -543,6 +547,40 @@ def test_parses_without_hierarchy_json(self) -> None: assert result.hierarchy is None assert result.manifest is not None + def test_parses_page_chunks_with_content_source_and_page_metadata(self) -> None: + manifest = _make_optimized_manifest() + manifest["statistics"]["total_chunks"] = 1 + manifest["statistics"]["text_chunks"] = 0 + manifest["statistics"]["image_chunks"] = 0 + manifest["statistics"]["table_chunks"] = 0 + manifest["statistics"]["page_chunks"] = 1 + chunks: List[Dict[str, Any]] = [ + { + "chunk_id": "page_chunk_1", + "type": "page", + "contentSource": "summary", + "path": "Default_Root/report.pdf-->pages/4-6", + "metadata": { + "page_nums": [4, 5, 6], + "summary": "The document explains refund eligibility.", + "entities": [{"text": "refund", "label": "topic"}], + }, + } + ] + zip_bytes = _build_zip(manifest, chunks=chunks) + + result = parseResultZip(zip_bytes, verify_checksum=False) + + assert result.manifest.statistics is not None + assert result.manifest.statistics.page_chunks == 1 + assert len(result.page_chunks) == 1 + page_chunk = result.page_chunks[0] + assert isinstance(page_chunk, PageChunk) + assert page_chunk.content_source == "summary" + assert page_chunk.content == "The document explains refund eligibility." + assert page_chunk.metadata.page_nums == [4, 5, 6] + assert page_chunk.metadata.entities == [{"text": "refund", "label": "topic"}] + # --------------------------------------------------------------------------- # Checksum verification @@ -576,9 +614,7 @@ def test_correct_checksum_passes(self) -> None: assert result.manifest is not None def test_wrong_checksum_raises_checksum_error(self) -> None: - manifest: Dict[str, Any] = _make_manifest( - checksum_value="wrong_checksum_value" - ) + manifest: Dict[str, Any] = _make_manifest(checksum_value="wrong_checksum_value") zip_bytes: bytes = _build_zip(manifest) with pytest.raises(ChecksumError): @@ -616,13 +652,9 @@ def test_missing_manifest_raises_error(self) -> None: def test_missing_chunks_returns_empty_chunks(self) -> None: """When chunks.json is missing, the result has no chunks.""" manifest: Dict[str, Any] = _make_manifest() - zip_bytes: bytes = _build_zip( - manifest, include_chunks=False - ) + zip_bytes: bytes = _build_zip(manifest, include_chunks=False) - result: ParseResult = parseResultZip( - zip_bytes, verify_checksum=False - ) + result: ParseResult = parseResultZip(zip_bytes, verify_checksum=False) assert result.chunks == [] @@ -782,6 +814,8 @@ def test_hierarchy_is_structured(self, parsed_real_result: ParseResult) -> None: # -- Raw ZIP -- - def test_raw_zip_is_preserved(self, real_zip_bytes: bytes, parsed_real_result: ParseResult) -> None: + def test_raw_zip_is_preserved( + self, real_zip_bytes: bytes, parsed_real_result: ParseResult + ) -> None: assert parsed_real_result.raw_zip is not None assert parsed_real_result.raw_zip == real_zip_bytes diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py index 7ce45d7..3632674 100644 --- a/tests/test_retrieval.py +++ b/tests/test_retrieval.py @@ -12,7 +12,7 @@ from tests.conftest import BASE_URL -RETRIEVAL_QUERY_URL: str = f"{BASE_URL}/v1/retrieval/query" +RETRIEVAL_QUERY_URL: str = f"{BASE_URL}/v2/retrieval/query" def _make_retrieval_response() -> Dict[str, Any]: @@ -26,7 +26,12 @@ def _make_retrieval_response() -> Dict[str, Any]: "failure_reason": "insufficient evidence", "decision_trace": [ {"phase": "discovery", "action": "select_documents", "selected": ["doc_123"]}, - {"phase": "terminal", "action": "complete", "stop_reason": "answer_done", "failure_reason": "insufficient evidence"}, + { + "phase": "terminal", + "action": "complete", + "stop_reason": "answer_done", + "failure_reason": "insufficient evidence", + }, ], "referenced_chunks": [ { @@ -41,7 +46,9 @@ def _make_retrieval_response() -> Dict[str, Any]: ], "results": [ { + "chunk_id": "chunk_001", "chunk_type": "text", + "content_source": "content", "content": "Annual plans may be refunded within 30 days.", "score": 1.0, "source": { @@ -135,9 +142,7 @@ def test_query_sends_request_and_returns_results(self, sync_client: Any) -> None assert response.results[0].source.document_id == "doc_123" assert response.results[0].source.source_file_name == "refund-policy.md" assert response.results[0].source.section_path == "Policies / Billing / Refunds" - assert response.answer_text == ( - "Annual plans may be refunded within 30 days of purchase." - ) + assert response.answer_text == ("Annual plans may be refunded within 30 days of purchase.") assert len(response.referenced_chunks) == 1 assert response.evidence_text == "Rendered retrieval evidence" assert response.stop_reason == "answer_done" @@ -146,9 +151,34 @@ def test_query_sends_request_and_returns_results(self, sync_client: Any) -> None assert response.referenced_chunks[0].chunk_type == "text" assert response.referenced_chunks[0].file_path is None assert not hasattr(response.results[0], "citation") - assert not hasattr(response.results[0], "chunk_id") + assert response.results[0].chunk_id == "chunk_001" + assert response.results[0].content_source == "content" assert not hasattr(response.results[0], "section_id") + @respx.mock + def test_query_sends_chunk_types( + self, + sync_client: Any, + ) -> None: + route = respx.post(RETRIEVAL_QUERY_URL).mock( + return_value=httpx.Response(200, json=_make_retrieval_response()) + ) + + response = sync_client.retrieval.query( + query="refund policy", + chunk_types=["page"], + data_type=7, + ) + + assert route.called + request_body: Dict[str, Any] = json.loads(route.calls[0].request.read()) + assert request_body == { + "query": "refund policy", + "data_type": 7, + "chunk_types": ["page"], + } + assert response.results[0].chunk_id == "chunk_001" + @respx.mock def test_query_omits_defaulted_optional_fields(self, sync_client: Any) -> None: route = respx.post(RETRIEVAL_QUERY_URL).mock( @@ -216,9 +246,7 @@ def test_agentic_response_fields(self, sync_client: Any) -> None: use_agentic=True, ) - assert response.answer_text == ( - "Annual plans may be refunded within 30 days of purchase." - ) + assert response.answer_text == ("Annual plans may be refunded within 30 days of purchase.") assert len(response.referenced_chunks) == 1 assert response.referenced_chunks[0].chunk_id == "chunk_001" assert response.referenced_chunks[0].document_id == "doc_123" @@ -226,9 +254,7 @@ def test_agentic_response_fields(self, sync_client: Any) -> None: assert response.referenced_chunks[0].section_path == "Policies / Billing / Refunds" assert response.referenced_chunks[0].file_path is None assert response.referenced_chunks[0].job_id == "job_123" - assert response.referenced_chunks[0].asset_url == ( - "https://example.com/assets/chunk_001" - ) + assert response.referenced_chunks[0].asset_url == ("https://example.com/assets/chunk_001") assert response.evidence_text == "Rendered retrieval evidence" assert response.stop_reason == "answer_done" assert response.failure_reason == "insufficient evidence" @@ -241,9 +267,7 @@ def test_agentic_response_fields(self, sync_client: Any) -> None: def test_legacy_response_without_agentic_fields(self, sync_client: Any) -> None: """Legacy-mode response defaults agentic fields to null and empty references.""" respx.post(RETRIEVAL_QUERY_URL).mock( - return_value=httpx.Response( - 200, json=_make_legacy_retrieval_response() - ) + return_value=httpx.Response(200, json=_make_legacy_retrieval_response()) ) response = sync_client.retrieval.query(query="refund policy") diff --git a/tests/test_retry.py b/tests/test_retry.py index 716e951..5ec2828 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -18,7 +18,7 @@ from tests.conftest import BASE_URL -JOBS_URL: str = f"{BASE_URL}/v1/jobs" +JOBS_URL: str = f"{BASE_URL}/v2/jobs" JOB_ID: str = "job_retry_test" GET_URL: str = f"{JOBS_URL}/{JOB_ID}" diff --git a/tests/test_upload.py b/tests/test_upload.py index 0b4c64b..0fb835a 100644 --- a/tests/test_upload.py +++ b/tests/test_upload.py @@ -129,7 +129,7 @@ def readable(self) -> bool: def seekable(self) -> bool: return False - def readinto(self, b: bytearray) -> int: + def readinto(self, b: Any) -> int: remaining: bytes = self._data[self._pos:] n: int = min(len(b), len(remaining)) b[:n] = remaining[:n]