Skip to content

Commit 73094d4

Browse files
author
Will
committed
feat: add document chunks resource methods
1 parent 87c249d commit 73094d4

7 files changed

Lines changed: 341 additions & 4 deletions

File tree

‎README.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,21 @@ update_job = client.jobs.create(
8585
document = client.documents.get(document_id)
8686
print(document.status)
8787

88+
chunks = client.documents.list_chunks(
89+
document_id,
90+
page=1,
91+
page_size=50,
92+
chunk_type="text",
93+
)
94+
print(chunks.pagination.total)
95+
if chunks.chunks:
96+
chunk = client.documents.get_chunk(
97+
document_id,
98+
chunks.chunks[0].id,
99+
include_asset_urls=True,
100+
)
101+
print(chunk.chunk.content)
102+
88103
client.documents.archive(document_id)
89104
```
90105

‎docs/usage.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -521,6 +521,22 @@ for document in document_list.documents:
521521
document = client.documents.get("doc_123")
522522
print(document.current_job_result_id)
523523

524+
chunks = client.documents.list_chunks(
525+
"doc_123",
526+
page=1,
527+
page_size=50,
528+
chunk_type="text",
529+
)
530+
for chunk in chunks.chunks:
531+
print(chunk.id, chunk.content)
532+
533+
image_chunk = client.documents.get_chunk(
534+
"doc_123",
535+
"dchk_123",
536+
include_asset_urls=True,
537+
)
538+
print(image_chunk.chunk.asset_url)
539+
524540
archived = client.documents.archive("doc_123")
525541
print(archived.status) # "archived"
526542
```

‎src/knowhere/__init__.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,15 @@
3535
)
3636
from knowhere._types import PollProgressCallback, UploadProgressCallback
3737
from knowhere._version import __version__
38-
from knowhere.types.document import Document, DocumentListResponse
38+
from knowhere.types.document import (
39+
Document,
40+
DocumentChunk,
41+
DocumentChunkListResponse,
42+
DocumentChunkPagination,
43+
DocumentChunkResponse,
44+
DocumentChunkType,
45+
DocumentListResponse,
46+
)
3947
from knowhere.types.job import Job, JobError, JobProgress, JobResult
4048
from knowhere.types.params import ParsingParams, WebhookConfig
4149
from knowhere.types.retrieval import (
@@ -98,6 +106,11 @@
98106
"JobResult",
99107
# Document types
100108
"Document",
109+
"DocumentChunk",
110+
"DocumentChunkListResponse",
111+
"DocumentChunkPagination",
112+
"DocumentChunkResponse",
113+
"DocumentChunkType",
101114
"DocumentListResponse",
102115
# Retrieval types
103116
"RetrievalChannel",

‎src/knowhere/resources/documents.py‎

Lines changed: 118 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,13 @@
55
from typing import Any, Dict, Optional
66

77
from knowhere.resources._base import AsyncAPIResource, SyncAPIResource
8-
from knowhere.types.document import Document, DocumentListResponse
8+
from knowhere.types.document import (
9+
Document,
10+
DocumentChunkListResponse,
11+
DocumentChunkResponse,
12+
DocumentChunkType,
13+
DocumentListResponse,
14+
)
915

1016

1117
class Documents(SyncAPIResource):
@@ -32,6 +38,49 @@ def get(self, document_id: str) -> Document:
3238
cast_to=Document,
3339
)
3440

41+
def list_chunks(
42+
self,
43+
document_id: str,
44+
*,
45+
page: int = 1,
46+
page_size: int = 50,
47+
chunk_type: Optional[DocumentChunkType] = None,
48+
include_asset_urls: bool = False,
49+
) -> DocumentChunkListResponse:
50+
"""List current-revision chunks for one canonical document."""
51+
params: Dict[str, Any] = _build_chunk_list_params(
52+
page=page,
53+
page_size=page_size,
54+
chunk_type=chunk_type,
55+
include_asset_urls=include_asset_urls,
56+
)
57+
58+
return self._request(
59+
"GET",
60+
f"v1/documents/{document_id}/chunks",
61+
params=params or None,
62+
cast_to=DocumentChunkListResponse,
63+
)
64+
65+
def get_chunk(
66+
self,
67+
document_id: str,
68+
document_chunk_id: str,
69+
*,
70+
include_asset_urls: bool = False,
71+
) -> DocumentChunkResponse:
72+
"""Get one current-revision chunk for one canonical document."""
73+
params: Dict[str, Any] = _build_chunk_get_params(
74+
include_asset_urls=include_asset_urls,
75+
)
76+
77+
return self._request(
78+
"GET",
79+
f"v1/documents/{document_id}/chunks/{document_chunk_id}",
80+
params=params or None,
81+
cast_to=DocumentChunkResponse,
82+
)
83+
3584
def archive(self, document_id: str) -> Document:
3685
"""Archive one canonical document by ID."""
3786
return self._request(
@@ -65,10 +114,78 @@ async def get(self, document_id: str) -> Document:
65114
cast_to=Document,
66115
)
67116

117+
async def list_chunks(
118+
self,
119+
document_id: str,
120+
*,
121+
page: int = 1,
122+
page_size: int = 50,
123+
chunk_type: Optional[DocumentChunkType] = None,
124+
include_asset_urls: bool = False,
125+
) -> DocumentChunkListResponse:
126+
"""List current-revision chunks for one canonical document."""
127+
params: Dict[str, Any] = _build_chunk_list_params(
128+
page=page,
129+
page_size=page_size,
130+
chunk_type=chunk_type,
131+
include_asset_urls=include_asset_urls,
132+
)
133+
134+
return await self._request(
135+
"GET",
136+
f"v1/documents/{document_id}/chunks",
137+
params=params or None,
138+
cast_to=DocumentChunkListResponse,
139+
)
140+
141+
async def get_chunk(
142+
self,
143+
document_id: str,
144+
document_chunk_id: str,
145+
*,
146+
include_asset_urls: bool = False,
147+
) -> DocumentChunkResponse:
148+
"""Get one current-revision chunk for one canonical document."""
149+
params: Dict[str, Any] = _build_chunk_get_params(
150+
include_asset_urls=include_asset_urls,
151+
)
152+
153+
return await self._request(
154+
"GET",
155+
f"v1/documents/{document_id}/chunks/{document_chunk_id}",
156+
params=params or None,
157+
cast_to=DocumentChunkResponse,
158+
)
159+
68160
async def archive(self, document_id: str) -> Document:
69161
"""Archive one canonical document by ID."""
70162
return await self._request(
71163
"POST",
72164
f"v1/documents/{document_id}/archive",
73165
cast_to=Document,
74166
)
167+
168+
169+
def _build_chunk_list_params(
170+
*,
171+
page: int,
172+
page_size: int,
173+
chunk_type: Optional[DocumentChunkType],
174+
include_asset_urls: bool,
175+
) -> Dict[str, Any]:
176+
params: Dict[str, Any] = {}
177+
if page != 1:
178+
params["page"] = page
179+
if page_size != 50:
180+
params["page_size"] = page_size
181+
if chunk_type is not None:
182+
params["chunk_type"] = chunk_type
183+
if include_asset_urls:
184+
params["include_asset_urls"] = True
185+
return params
186+
187+
188+
def _build_chunk_get_params(*, include_asset_urls: bool) -> Dict[str, Any]:
189+
if not include_asset_urls:
190+
return {}
191+
return {"include_asset_urls": True}

‎src/knowhere/types/__init__.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,15 @@
22

33
from __future__ import annotations
44

5-
from knowhere.types.document import Document, DocumentListResponse
5+
from knowhere.types.document import (
6+
Document,
7+
DocumentChunk,
8+
DocumentChunkListResponse,
9+
DocumentChunkPagination,
10+
DocumentChunkResponse,
11+
DocumentChunkType,
12+
DocumentListResponse,
13+
)
614
from knowhere.types.job import Job, JobError, JobResult
715
from knowhere.types.params import ParsingParams, WebhookConfig
816
from knowhere.types.retrieval import (
@@ -39,6 +47,11 @@
3947
"JobResult",
4048
# document
4149
"Document",
50+
"DocumentChunk",
51+
"DocumentChunkListResponse",
52+
"DocumentChunkPagination",
53+
"DocumentChunkResponse",
54+
"DocumentChunkType",
4255
"DocumentListResponse",
4356
# retrieval
4457
"RetrievalChannel",

‎src/knowhere/types/document.py‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from __future__ import annotations
44

55
from datetime import datetime
6-
from typing import Optional
6+
from typing import Any, Dict, Literal, Optional
77

88
from pydantic import BaseModel
99

@@ -26,3 +26,53 @@ class DocumentListResponse(BaseModel):
2626

2727
namespace: str
2828
documents: list[Document]
29+
30+
31+
DocumentChunkType = Literal["text", "image", "table"]
32+
33+
34+
class DocumentChunkPagination(BaseModel):
35+
"""Pagination metadata returned by document chunk list endpoints."""
36+
37+
page: int
38+
page_size: int
39+
total: int
40+
total_pages: int
41+
42+
43+
class DocumentChunk(BaseModel):
44+
"""One current-revision document chunk."""
45+
46+
id: str
47+
chunk_id: str
48+
chunk_type: DocumentChunkType
49+
content: Optional[str] = None
50+
section_id: Optional[str] = None
51+
section_path: Optional[str] = None
52+
source_chunk_path: Optional[str] = None
53+
file_path: Optional[str] = None
54+
sort_order: int
55+
metadata: Dict[str, Any]
56+
asset_url: Optional[str] = None
57+
created_at: Optional[datetime] = None
58+
59+
60+
class DocumentChunkListResponse(BaseModel):
61+
"""Response from ``GET /v1/documents/{document_id}/chunks``."""
62+
63+
document_id: str
64+
namespace: str
65+
job_result_id: Optional[str] = None
66+
job_id: Optional[str] = None
67+
chunks: list[DocumentChunk]
68+
pagination: DocumentChunkPagination
69+
70+
71+
class DocumentChunkResponse(BaseModel):
72+
"""Response from ``GET /v1/documents/{document_id}/chunks/{chunk_id}``."""
73+
74+
document_id: str
75+
namespace: str
76+
job_result_id: Optional[str] = None
77+
job_id: Optional[str] = None
78+
chunk: DocumentChunk

0 commit comments

Comments
 (0)