knowledge2 0.4.0__py3-none-any.whl

sdk/async_resources/documents.py

@@ -0,0 +1,592 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from typing import Any
+
+from sdk._async_paging import AsyncPager
+from sdk._paging import Page
+from sdk._request_options import RequestOptions
+from sdk._validation import require_str
+from sdk.async_resources._mixin_base import AsyncRequesterMixin
+from sdk.errors import ConfirmationRequiredError
+from sdk.types import (
+    ChunkingConfig,
+    DocumentBatchUploadResponse,
+    DocumentCreateResponse,
+    DocumentDeleteResponse,
+    DocumentDetailResponse,
+    DocumentManifestIngestResponse,
+    DocumentUrlIngestResponse,
+)
+
+
+class AsyncDocumentsMixin(AsyncRequesterMixin):
+    async def upload_document(
+        self,
+        corpus_id: str,
+        *,
+        file_path: str | None = None,
+        file_bytes: bytes | None = None,
+        filename: str | None = None,
+        raw_text: str | None = None,
+        source_uri: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        auto_index: bool | None = None,
+        chunk_strategy: str | None = None,
+        chunking: ChunkingConfig | None = None,
+        idempotency_key: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentCreateResponse:
+        """Upload a document to a corpus.
+
+        Args:
+            corpus_id: Target corpus ID.
+            file_path: Path to file to upload.
+            file_bytes: Raw file bytes to upload.
+            filename: Filename when using file_bytes.
+            raw_text: Raw text content to upload.
+            source_uri: Optional source URI for the document.
+            metadata: Optional document metadata.
+            auto_index: Whether to auto-index after ingestion.
+            chunk_strategy: Deprecated - use chunking instead.
+            chunking: Chunking configuration (strategy, chunk_size, overlap, etc.)
+            idempotency_key: Optional key for idempotent requests.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        if file_path and (file_bytes or raw_text):
+            raise ValueError("file_path cannot be combined with file_bytes or raw_text")
+        if file_bytes and raw_text:
+            raise ValueError("file_bytes cannot be combined with raw_text")
+        headers = self._idempotency_headers(idempotency_key)
+        if file_path:
+            form: dict[str, str] = {}
+            if source_uri is not None:
+                form["source_uri"] = source_uri
+            if metadata is not None:
+                form["metadata"] = json.dumps(metadata)
+            if auto_index is not None:
+                form["auto_index"] = str(bool(auto_index)).lower()
+            if chunking is not None:
+                form["chunking"] = json.dumps(chunking)
+            elif chunk_strategy is not None:
+                form["chunk_strategy"] = chunk_strategy
+            with open(file_path, "rb") as handle:
+                files = {"file": (os.path.basename(file_path), handle)}
+                data = await self._request(
+                    "POST",
+                    f"/v1/corpora/{corpus_id}/documents",
+                    data=form,
+                    files=files,
+                    headers=headers,
+                    request_options=request_options,
+                )
+            return self._maybe_validate(data, "DocumentCreateResponse")
+        if file_bytes is not None:
+            if not filename:
+                raise ValueError("filename is required when using file_bytes")
+            form_data: dict[str, str] = {}
+            if source_uri is not None:
+                form_data["source_uri"] = source_uri
+            if metadata is not None:
+                form_data["metadata"] = json.dumps(metadata)
+            if auto_index is not None:
+                form_data["auto_index"] = str(bool(auto_index)).lower()
+            if chunking is not None:
+                form_data["chunking"] = json.dumps(chunking)
+            elif chunk_strategy is not None:
+                form_data["chunk_strategy"] = chunk_strategy
+            file_payload: dict[str, Any] = {"file": (filename, file_bytes)}
+            data = await self._request(
+                "POST",
+                f"/v1/corpora/{corpus_id}/documents",
+                data=form_data,
+                files=file_payload,
+                headers=headers,
+                request_options=request_options,
+            )
+            return self._maybe_validate(data, "DocumentCreateResponse")
+        if raw_text is None:
+            raise ValueError("raw_text is required when no file is provided")
+        payload: dict[str, Any] = {}
+        if raw_text is not None:
+            payload["raw_text"] = raw_text
+        if source_uri is not None:
+            payload["source_uri"] = source_uri
+        if metadata is not None:
+            payload["metadata"] = metadata
+        if auto_index is not None:
+            payload["auto_index"] = auto_index
+        if chunking is not None:
+            payload["chunking"] = chunking
+        elif chunk_strategy is not None:
+            payload["chunk_strategy"] = chunk_strategy
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/documents",
+            json=payload,
+            headers=headers,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "DocumentCreateResponse")
+
+    async def upload_documents_batch(
+        self,
+        corpus_id: str,
+        documents: list[dict[str, Any]],
+        idempotency_key: str | None = None,
+        *,
+        auto_index: bool | None = None,
+        chunk_strategy: str | None = None,
+        chunking: ChunkingConfig | None = None,
+        wait: bool = True,
+        poll_s: int = 5,
+        timeout_s: float | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentBatchUploadResponse:
+        """Upload multiple documents as raw text in a batch.
+
+        Args:
+            corpus_id: Target corpus ID.
+            documents: List of document dicts with raw_text, source_uri, metadata.
+            idempotency_key: Optional key for idempotent requests.
+            auto_index: Whether to auto-index after ingestion.
+            chunk_strategy: Deprecated - use chunking instead.
+            chunking: Chunking configuration (strategy, chunk_size, overlap, etc.)
+            wait: If True, wait for the batch job to complete.
+            poll_s: Polling interval when waiting.
+            timeout_s: Maximum seconds to wait for job completion.
+
+        Returns:
+            Response with ``doc_ids`` (list of created document IDs),
+            ``job_id``, and ``count``.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"documents": documents}
+        if auto_index is not None:
+            payload["auto_index"] = auto_index
+        if chunking is not None:
+            payload["chunking"] = chunking
+        elif chunk_strategy is not None:
+            payload["chunk_strategy"] = chunk_strategy
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/documents:batch",
+            json=payload,
+            headers=headers,
+            request_options=request_options,
+        )
+        if wait:
+            job_id = data.get("job_id")
+            if job_id:
+                await self._wait_for_job(job_id, poll_s=poll_s, timeout_s=timeout_s)
+        return self._maybe_validate(data, "DocumentBatchUploadResponse")
+
+    async def upload_files_batch(
+        self,
+        corpus_id: str,
+        files: list[tuple[str, bytes]],
+        idempotency_key: str | None = None,
+        *,
+        auto_index: bool | None = None,
+        chunk_strategy: str | None = None,
+        chunking: ChunkingConfig | None = None,
+        wait: bool = True,
+        poll_s: int = 5,
+        timeout_s: float | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentBatchUploadResponse:
+        """Upload multiple files in a single multipart request.
+
+        Args:
+            corpus_id: Target corpus ID.
+            files: List of (filename, content_bytes) tuples.
+            idempotency_key: Optional key for idempotent requests.
+            auto_index: Whether to auto-index after ingestion.
+            chunk_strategy: Deprecated - use chunking instead.
+            chunking: Chunking configuration (strategy, chunk_size, overlap, etc.)
+            wait: If True, wait for the batch job to complete.
+            poll_s: Polling interval when waiting.
+            timeout_s: Maximum seconds to wait for job completion.
+
+        Returns:
+            Response with job_id, doc_ids, and count.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        headers = self._idempotency_headers(idempotency_key)
+
+        files_list = [("files", (filename, content)) for filename, content in files]
+
+        form_data: dict[str, Any] = {}
+        if auto_index is not None:
+            form_data["auto_index"] = str(auto_index).lower()
+        if chunking is not None:
+            form_data["chunking"] = json.dumps(chunking)
+        elif chunk_strategy is not None:
+            form_data["chunk_strategy"] = chunk_strategy
+
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/documents:upload_batch",
+            data=form_data if form_data else None,
+            files=files_list,
+            headers=headers,
+            request_options=request_options,
+        )
+        if wait:
+            job_id = data.get("job_id")
+            if job_id:
+                await self._wait_for_job(job_id, poll_s=poll_s, timeout_s=timeout_s)
+        return self._maybe_validate(data, "DocumentBatchUploadResponse")
+
+    async def upload_documents_parallel(
+        self,
+        corpus_id: str,
+        file_paths: list[str],
+        *,
+        max_workers: int = 8,
+        auto_index: bool | None = None,
+        chunking: ChunkingConfig | None = None,
+        metadata: dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> list[DocumentCreateResponse]:
+        """Upload multiple files concurrently using asyncio.gather.
+
+        Each file is uploaded as a separate HTTP request via
+        :meth:`upload_document`.  For large batches this is significantly
+        faster than uploading sequentially.
+
+        Unlike :meth:`upload_files_batch` (which sends a single multipart
+        request for server-side batching), this method issues concurrent
+        individual uploads for client-side parallelism.
+
+        .. warning:: **All-or-nothing semantics.** On partial failure an
+           :class:`ExceptionGroup` is raised and **no** successful results are
+           returned.  Callers who need partial-failure recovery should upload
+           files individually via :meth:`upload_document` and handle errors
+           per file.
+
+        Args:
+            corpus_id: Target corpus ID.
+            file_paths: List of local file paths to upload.
+            max_workers: Maximum number of concurrent upload coroutines
+                (default 8, must be >= 1).
+            auto_index: Whether to auto-index after ingestion.
+            chunking: Chunking configuration applied to each upload.
+            metadata: Optional metadata applied to every document.
+
+        Returns:
+            List of upload responses (one per file, in input order).
+
+        Raises:
+            ExceptionGroup: If one or more uploads fail, containing all
+                individual exceptions from failed uploads.  Successful
+                results from other uploads are discarded.
+            ValueError: If *max_workers* is less than 1 or greater than 256.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        if max_workers < 1:
+            raise ValueError(f"max_workers must be >= 1, got {max_workers}")
+        if max_workers > 256:
+            raise ValueError(f"max_workers must be <= 256, got {max_workers}")
+        if not file_paths:
+            return []
+
+        semaphore = asyncio.Semaphore(max_workers)
+
+        async def _upload(fp: str) -> DocumentCreateResponse:
+            async with semaphore:
+                return await self.upload_document(
+                    corpus_id,
+                    file_path=fp,
+                    auto_index=auto_index,
+                    chunking=chunking,
+                    metadata=metadata,
+                    request_options=request_options,
+                )
+
+        tasks = [_upload(fp) for fp in file_paths]
+        settled = await asyncio.gather(*tasks, return_exceptions=True)
+
+        results: list[DocumentCreateResponse] = []
+        errors: list[Exception] = []
+        for outcome in settled:
+            if isinstance(outcome, Exception):
+                errors.append(outcome)
+            elif isinstance(outcome, BaseException):
+                err = RuntimeError(str(outcome))
+                err.__cause__ = outcome
+                errors.append(err)
+            else:
+                results.append(outcome)
+
+        if errors:
+            raise ExceptionGroup(f"{len(errors)} of {len(file_paths)} uploads failed", errors)
+        return results
+
+    async def ingest_urls(
+        self,
+        corpus_id: str,
+        urls: list[dict[str, Any]],
+        idempotency_key: str | None = None,
+        *,
+        auto_index: bool | None = None,
+        chunk_strategy: str | None = None,
+        chunking: ChunkingConfig | None = None,
+        wait: bool = True,
+        poll_s: int = 5,
+        timeout_s: float | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentUrlIngestResponse:
+        """Ingest documents from URLs.
+
+        Args:
+            corpus_id: Target corpus ID.
+            urls: List of URL dicts with url, title, tags, metadata.
+            idempotency_key: Optional key for idempotent requests.
+            auto_index: Whether to auto-index after ingestion.
+            chunk_strategy: Deprecated - use chunking instead.
+            chunking: Chunking configuration (strategy, chunk_size, overlap, etc.)
+            wait: If True, wait for the batch job to complete.
+            poll_s: Polling interval when waiting.
+            timeout_s: Maximum seconds to wait for job completion.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"urls": urls}
+        if auto_index is not None:
+            payload["auto_index"] = auto_index
+        if chunking is not None:
+            payload["chunking"] = chunking
+        elif chunk_strategy is not None:
+            payload["chunk_strategy"] = chunk_strategy
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/documents:ingest_urls",
+            json=payload,
+            headers=headers,
+            request_options=request_options,
+        )
+        if wait:
+            job_id = data.get("job_id")
+            if job_id:
+                await self._wait_for_job(job_id, poll_s=poll_s, timeout_s=timeout_s)
+        return self._maybe_validate(data, "DocumentUrlIngestResponse")
+
+    async def ingest_manifest(
+        self,
+        corpus_id: str,
+        manifest_uri: str,
+        max_documents: int | None = None,
+        idempotency_key: str | None = None,
+        *,
+        auto_index: bool | None = None,
+        chunk_strategy: str | None = None,
+        chunking: ChunkingConfig | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentManifestIngestResponse:
+        """Ingest documents from a manifest file.
+
+        Args:
+            corpus_id: Target corpus ID.
+            manifest_uri: URI to manifest file (S3, HTTP, local).
+            max_documents: Optional limit on documents to ingest.
+            idempotency_key: Optional key for idempotent requests.
+            auto_index: Whether to auto-index after ingestion.
+            chunk_strategy: Deprecated - use chunking instead.
+            chunking: Chunking configuration (strategy, chunk_size, overlap, etc.)
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"manifest_uri": manifest_uri}
+        if max_documents is not None:
+            payload["max_documents"] = max_documents
+        if auto_index is not None:
+            payload["auto_index"] = auto_index
+        if chunking is not None:
+            payload["chunking"] = chunking
+        elif chunk_strategy is not None:
+            payload["chunk_strategy"] = chunk_strategy
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/documents:ingest_manifest",
+            json=payload,
+            headers=headers,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "DocumentManifestIngestResponse")
+
+    async def list_documents(
+        self,
+        corpus_id: str,
+        *,
+        limit: int = 100,
+        offset: int = 0,
+        q: str | None = None,
+        status: str | None = None,
+        source: str | None = None,
+        tag: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> Page[dict[str, Any]]:
+        corpus_id = require_str(corpus_id, "corpus_id")
+        params: dict[str, Any] = {}
+        if q is not None:
+            params["q"] = q
+        if status is not None:
+            params["status"] = status
+        if source is not None:
+            params["source"] = source
+        if tag is not None:
+            params["tag"] = tag
+        return await self._list_page(
+            "GET",
+            f"/v1/corpora/{corpus_id}/documents",
+            items_key="documents",
+            params=params or None,
+            limit=limit,
+            offset=offset,
+        )
+
+    def iter_documents(
+        self,
+        corpus_id: str,
+        *,
+        limit: int = 100,
+        q: str | None = None,
+        status: str | None = None,
+        source: str | None = None,
+        tag: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> AsyncPager[dict[str, Any]]:
+        """Lazily paginate documents, yielding individual document items."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        params: dict[str, Any] = {}
+        if q is not None:
+            params["q"] = q
+        if status is not None:
+            params["status"] = status
+        if source is not None:
+            params["source"] = source
+        if tag is not None:
+            params["tag"] = tag
+        return self._paginate(
+            "GET",
+            f"/v1/corpora/{corpus_id}/documents",
+            items_key="documents",
+            params=params if params else None,
+            limit=limit,
+        )
+
+    async def get_document(
+        self,
+        doc_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentDetailResponse:
+        doc_id = require_str(doc_id, "doc_id")
+        data = await self._request(
+            "GET", f"/v1/documents/{doc_id}", request_options=request_options
+        )
+        return self._maybe_validate(data, "DocumentDetailResponse")
+
+    async def update_document_metadata(
+        self,
+        doc_id: str,
+        metadata: dict[str, Any],
+        request_options: RequestOptions | None = None,
+    ) -> dict[str, Any]:
+        """Update customer metadata on a document using merge semantics.
+
+        Keys with non-empty values are added or updated.
+        Keys with empty string or None values are removed.
+        Keys not in the request are left unchanged.
+
+        Args:
+            doc_id: Document ID to update
+            metadata: Dict of metadata updates to apply
+
+        Returns:
+            Updated metadata dict with custom_metadata and system_metadata
+        """
+        doc_id = require_str(doc_id, "doc_id")
+        response = await self._request(
+            "PATCH",
+            f"/v1/documents/{doc_id}/metadata",
+            json=metadata,
+            request_options=request_options,
+        )
+        return response
+
+    async def delete_document(
+        self,
+        corpus_id: str,
+        doc_id: str,
+        *,
+        confirm: bool = False,
+        reindex: bool = False,
+        request_options: RequestOptions | None = None,
+    ) -> DocumentDeleteResponse:
+        """Delete a document from a corpus.
+
+        This is an irreversible operation.  You must pass ``confirm=True``
+        to acknowledge this and proceed.
+
+        Args:
+            corpus_id: The corpus containing the document.
+            doc_id: Unique identifier of the document to delete.
+            confirm: Safety guard -- must be ``True`` to execute the
+                deletion.  Raises ``ConfirmationRequiredError`` when ``False``.
+            reindex: If ``True``, trigger a re-index of the corpus after
+                deletion.
+
+        Returns:
+            Confirmation of the deletion.
+
+        Raises:
+            ConfirmationRequiredError: If *confirm* is not ``True``.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        doc_id = require_str(doc_id, "doc_id")
+        if not confirm:
+            raise ConfirmationRequiredError("document", doc_id)
+        data = await self._request(
+            "DELETE",
+            f"/v1/corpora/{corpus_id}/documents/{doc_id}",
+            params={"reindex": reindex},
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "DocumentDeleteResponse")
+
+    async def list_chunks(
+        self,
+        corpus_id: str,
+        limit: int = 100,
+        offset: int = 0,
+        request_options: RequestOptions | None = None,
+    ) -> Page[dict[str, Any]]:
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return await self._list_page(
+            "GET",
+            f"/v1/corpora/{corpus_id}/chunks",
+            items_key="chunks",
+            limit=limit,
+            offset=offset,
+        )
+
+    def iter_chunks(
+        self,
+        corpus_id: str,
+        *,
+        limit: int = 100,
+        request_options: RequestOptions | None = None,
+    ) -> AsyncPager[dict[str, Any]]:
+        """Lazily paginate chunks, yielding individual chunk items."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return self._paginate(
+            "GET",
+            f"/v1/corpora/{corpus_id}/chunks",
+            items_key="chunks",
+            limit=limit,
+        )