knowledge2 0.4.0__py3-none-any.whl

sdk/async_resources/search.py

@@ -0,0 +1,262 @@
+"""Async search resource mixin for the Knowledge2 SDK."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sdk._request_options import RequestOptions
+from sdk._validation import require_str
+from sdk.async_resources._mixin_base import AsyncRequesterMixin
+from sdk.types import (
+    EmbeddingModelListResponse,
+    EmbeddingsResponse,
+    FeedbackResponse,
+    SearchBatchResponse,
+    SearchGenerateResponse,
+    SearchResponse,
+)
+from sdk.types.search import (
+    SearchGenerationConfig,
+    SearchHybridConfig,
+    SearchRerankConfig,
+    SearchReturnConfig,
+)
+
+
+class AsyncSearchMixin(AsyncRequesterMixin):
+    async def search(
+        self,
+        corpus_id: str,
+        query: str,
+        *,
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+        hybrid: SearchHybridConfig | dict[str, Any] | None = None,
+        rerank: SearchRerankConfig | dict[str, Any] | None = None,
+        return_config: SearchReturnConfig | dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> SearchResponse:
+        """Search a corpus for relevant chunks.
+
+        Args:
+            corpus_id: The corpus to search.
+            query: The search query text.
+            top_k: Maximum number of results to return.
+            filters: Optional metadata filters to narrow results.
+            hybrid: Hybrid search configuration (dense/sparse weighting).
+            rerank: Reranking configuration applied after initial retrieval.
+            return_config: Controls which fields are included in the
+                response (e.g. text, metadata, provenance).
+
+        Returns:
+            Search results with scored chunks, metadata, and text.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"query": query, "top_k": top_k}
+        if filters is not None:
+            payload["filters"] = filters
+        if hybrid is not None:
+            payload["hybrid"] = hybrid
+        if rerank is not None:
+            payload["rerank"] = rerank
+        if return_config is not None:
+            payload["return"] = return_config
+        data = await self._request(
+            "POST", f"/v1/corpora/{corpus_id}/search", json=payload, request_options=request_options
+        )
+        return self._maybe_validate(data, "SearchResponse")
+
+    async def search_batch(
+        self,
+        corpus_id: str,
+        queries: list[str],
+        *,
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+        hybrid: SearchHybridConfig | dict[str, Any] | None = None,
+        rerank: SearchRerankConfig | dict[str, Any] | None = None,
+        return_config: SearchReturnConfig | dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> SearchBatchResponse:
+        """Execute multiple search queries against a corpus in a single request.
+
+        Args:
+            corpus_id: The corpus to search.
+            queries: List of search query texts.
+            top_k: Maximum number of results to return per query.
+            filters: Optional metadata filters applied to all queries.
+            hybrid: Hybrid search configuration (dense/sparse weighting).
+            rerank: Reranking configuration applied after initial retrieval.
+            return_config: Controls which fields are included in each
+                query's results.
+
+        Returns:
+            Batch search results -- one result set per input query.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"queries": queries, "top_k": top_k}
+        if filters is not None:
+            payload["filters"] = filters
+        if hybrid is not None:
+            payload["hybrid"] = hybrid
+        if rerank is not None:
+            payload["rerank"] = rerank
+        if return_config is not None:
+            payload["return"] = return_config
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/search:batch",
+            json=payload,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "SearchBatchResponse")
+
+    async def search_generate(
+        self,
+        corpus_id: str,
+        query: str,
+        *,
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+        hybrid: SearchHybridConfig | dict[str, Any] | None = None,
+        rerank: SearchRerankConfig | dict[str, Any] | None = None,
+        return_config: SearchReturnConfig | dict[str, Any] | None = None,
+        generation: SearchGenerationConfig | dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> SearchGenerateResponse:
+        """Search a corpus and generate an LLM answer grounded in the results.
+
+        Args:
+            corpus_id: The corpus to search.
+            query: The search query text.
+            top_k: Maximum number of retrieval results fed to the generator.
+            filters: Optional metadata filters to narrow retrieval results.
+            hybrid: Hybrid search configuration (dense/sparse weighting).
+            rerank: Reranking configuration applied after initial retrieval.
+            return_config: Controls which fields are included in the
+                retrieval results.
+            generation: LLM generation settings (model, temperature,
+                system prompt, etc.).
+
+        Returns:
+            The generated answer together with the supporting retrieval
+            results.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {"query": query, "top_k": top_k}
+        if filters is not None:
+            payload["filters"] = filters
+        if hybrid is not None:
+            payload["hybrid"] = hybrid
+        if rerank is not None:
+            payload["rerank"] = rerank
+        if return_config is not None:
+            payload["return"] = return_config
+        if generation is not None:
+            payload["generation"] = generation
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/search:generate",
+            json=payload,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "SearchGenerateResponse")
+
+    async def embeddings(
+        self,
+        model: str,
+        inputs: list[str],
+        embed_type: str = "query",
+        request_options: RequestOptions | None = None,
+    ) -> EmbeddingsResponse:
+        """Generate vector embeddings for the given texts.
+
+        Args:
+            model: Name or ID of the embedding model to use.
+            inputs: List of text strings to embed.
+            embed_type: Embedding type -- ``"query"`` for search queries
+                or ``"document"`` for corpus documents.
+
+        Returns:
+            Embedding vectors for each input text.
+
+        Raises:
+            Knowledge2Error: If the API request fails.
+        """
+        payload = {"model": model, "input": inputs, "type": embed_type}
+        data = await self._request(
+            "POST", "/v1/embeddings", json=payload, request_options=request_options
+        )
+        return self._maybe_validate(data, "EmbeddingsResponse")
+
+    async def list_embedding_models(
+        self,
+        request_options: RequestOptions | None = None,
+    ) -> EmbeddingModelListResponse:
+        """List available embedding models.
+
+        Returns the set of embedding models that can be used with the
+        ``embeddings`` method, indicating which model is the default.
+
+        Returns:
+            The list of available embedding models with default flag.
+
+        Raises:
+            Knowledge2Error: If the API request fails.
+        """
+        data = await self._request("GET", "/v1/embeddings/models", request_options=request_options)
+        return self._maybe_validate(data, "EmbeddingModelListResponse")
+
+    async def create_feedback(
+        self,
+        corpus_id: str,
+        query: str,
+        *,
+        clicked_chunk_ids: list[str] | None = None,
+        rating: int | None = None,
+        abstained: bool = False,
+        request_options: RequestOptions | None = None,
+    ) -> FeedbackResponse:
+        """Submit relevance feedback for a search query.
+
+        Args:
+            corpus_id: The corpus the query was executed against.
+            query: The original search query text.
+            clicked_chunk_ids: IDs of chunks the user found relevant.
+            rating: Overall relevance rating for the result set.
+            abstained: If ``True``, indicates the user chose not to
+                rate the results.
+
+        Returns:
+            Confirmation that the feedback was recorded.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        payload: dict[str, Any] = {
+            "query": query,
+            "clicked_chunk_ids": clicked_chunk_ids,
+            "rating": rating,
+            "abstained": abstained,
+        }
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/feedback",
+            json=payload,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "FeedbackResponse")

sdk/async_resources/training.py

@@ -0,0 +1,357 @@
+"""Async training and tuning resource mixin for the Knowledge2 SDK."""
+
+from __future__ import annotations
+
+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.types import (
+    CancelTuningRunResponse,
+    EvalRunDetailResponse,
+    PromoteResponse,
+    TrainingDataBuildResponse,
+    TuningRunBuildResponse,
+    TuningRunDetailResponse,
+    TuningRunLogsResponse,
+    TuningRunResponse,
+)
+
+
+class AsyncTrainingMixin(AsyncRequesterMixin):
+    async def build_training_data(
+        self,
+        corpus_id: str,
+        idempotency_key: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> TrainingDataBuildResponse:
+        """Trigger a training data build job for a corpus.
+
+        Args:
+            corpus_id: The corpus to build training data from.
+            idempotency_key: Optional idempotency key to prevent
+                duplicate builds.
+
+        Returns:
+            The training data build response, including the background
+            job ID.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            ConflictError: If a duplicate idempotency key is detected.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/training-data:build",
+            json={},
+            headers=headers,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "TrainingDataBuildResponse")
+
+    async def list_training_data(
+        self,
+        corpus_id: str,
+        limit: int = 100,
+        offset: int = 0,
+        request_options: RequestOptions | None = None,
+    ) -> Page[dict[str, Any]]:
+        """List training datasets generated for a corpus.
+
+        Args:
+            corpus_id: The corpus whose training data to list.
+            limit: Maximum number of datasets to return per page.
+            offset: Number of datasets to skip for pagination.
+
+        Returns:
+            A Page containing training dataset records with pagination metadata.
+
+        Raises:
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return await self._list_page(
+            "GET",
+            f"/v1/corpora/{corpus_id}/training-data",
+            items_key="datasets",
+            limit=limit,
+            offset=offset,
+        )
+
+    def iter_training_data(
+        self,
+        corpus_id: str,
+        *,
+        limit: int = 100,
+        request_options: RequestOptions | None = None,
+    ) -> AsyncPager[dict[str, Any]]:
+        """Iterate over training data, automatically paginating.
+
+        Args:
+            corpus_id: The corpus whose training data to iterate.
+            limit: Page size used for each underlying API request.
+
+        Yields:
+            Individual training dataset dicts.
+
+        Raises:
+            Knowledge2Error: If any underlying API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return self._paginate(
+            "GET",
+            f"/v1/corpora/{corpus_id}/training-data",
+            items_key="datasets",
+            limit=limit,
+        )
+
+    async def list_tuning_runs(
+        self,
+        corpus_id: str,
+        limit: int = 100,
+        offset: int = 0,
+        request_options: RequestOptions | None = None,
+    ) -> Page[dict[str, Any]]:
+        """List tuning runs for a corpus.
+
+        Args:
+            corpus_id: The corpus whose tuning runs to list.
+            limit: Maximum number of tuning runs to return per page.
+            offset: Number of tuning runs to skip for pagination.
+
+        Returns:
+            A Page containing tuning run records with pagination metadata.
+
+        Raises:
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return await self._list_page(
+            "GET",
+            f"/v1/corpora/{corpus_id}/tuning-runs",
+            items_key="runs",
+            limit=limit,
+            offset=offset,
+        )
+
+    def iter_tuning_runs(
+        self,
+        corpus_id: str,
+        *,
+        limit: int = 100,
+        request_options: RequestOptions | None = None,
+    ) -> AsyncPager[dict[str, Any]]:
+        """Iterate over tuning runs, automatically paginating.
+
+        Args:
+            corpus_id: The corpus whose tuning runs to iterate.
+            limit: Page size used for each underlying API request.
+
+        Yields:
+            Individual tuning run dicts.
+
+        Raises:
+            Knowledge2Error: If any underlying API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        return self._paginate(
+            "GET",
+            f"/v1/corpora/{corpus_id}/tuning-runs",
+            items_key="runs",
+            limit=limit,
+        )
+
+    async def create_tuning_run(
+        self,
+        corpus_id: str,
+        idempotency_key: str | None = None,
+        *,
+        request_options: RequestOptions | None = None,
+    ) -> TuningRunResponse:
+        """Create a tuning run for the corpus.
+
+        Args:
+            corpus_id: The corpus to tune.
+            idempotency_key: Optional idempotency key.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/tuning-runs",
+            json={},
+            headers=headers,
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "TuningRunResponse")
+
+    async def build_and_start_tuning_run(
+        self,
+        corpus_id: str,
+        idempotency_key: str | None = None,
+        *,
+        wait: bool = True,
+        poll_s: int = 5,
+        request_options: RequestOptions | None = None,
+    ) -> TuningRunBuildResponse:
+        """Build training data and start a tuning run in one step.
+
+        Args:
+            corpus_id: The corpus to tune.
+            idempotency_key: Optional idempotency key to prevent
+                duplicate runs.
+            wait: If ``True`` (default), block until the build job
+                completes.
+            poll_s: Polling interval in seconds when *wait* is ``True``.
+
+        Returns:
+            The tuning run build response, including the background
+            job ID.
+
+        Raises:
+            NotFoundError: If the corpus does not exist.
+            ConflictError: If a duplicate idempotency key is detected.
+            Knowledge2Error: If the API request fails.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        headers = self._idempotency_headers(idempotency_key)
+        data = await self._request(
+            "POST",
+            f"/v1/corpora/{corpus_id}/tuning-runs:build",
+            json={},
+            headers=headers,
+            request_options=request_options,
+        )
+        if wait:
+            job_id = data.get("build_job_id") or data.get("job_id")
+            if job_id:
+                await self._wait_for_job(job_id, poll_s=poll_s)
+        return self._maybe_validate(data, "TuningRunBuildResponse")
+
+    async def get_tuning_run(
+        self,
+        run_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> TuningRunDetailResponse:
+        """Retrieve details of a tuning run.
+
+        Args:
+            run_id: Unique identifier of the tuning run.
+
+        Returns:
+            Detailed tuning run information including status, metrics,
+            and configuration.
+
+        Raises:
+            NotFoundError: If the tuning run does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        run_id = require_str(run_id, "run_id")
+        data = await self._request(
+            "GET", f"/v1/tuning-runs/{run_id}", request_options=request_options
+        )
+        return self._maybe_validate(data, "TuningRunDetailResponse")
+
+    async def get_tuning_run_logs(
+        self,
+        run_id: str,
+        tail: int = 200,
+        request_options: RequestOptions | None = None,
+    ) -> TuningRunLogsResponse:
+        """Retrieve log output from a tuning run.
+
+        Args:
+            run_id: Unique identifier of the tuning run.
+            tail: Number of most recent log lines to return.
+
+        Returns:
+            The requested log lines from the tuning run.
+
+        Raises:
+            NotFoundError: If the tuning run does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        run_id = require_str(run_id, "run_id")
+        data = await self._request(
+            "GET",
+            f"/v1/tuning-runs/{run_id}/logs",
+            params={"tail": tail},
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "TuningRunLogsResponse")
+
+    async def cancel_tuning_run(
+        self,
+        run_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> CancelTuningRunResponse:
+        """Cancel a running or pending tuning run.
+
+        Args:
+            run_id: Unique identifier of the tuning run to cancel.
+
+        Returns:
+            Confirmation of the cancellation.
+
+        Raises:
+            NotFoundError: If the tuning run does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        run_id = require_str(run_id, "run_id")
+        data = await self._request(
+            "POST", f"/v1/tuning-runs/{run_id}:cancel", request_options=request_options
+        )
+        return self._maybe_validate(data, "CancelTuningRunResponse")
+
+    async def promote_tuning_run(
+        self,
+        run_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> PromoteResponse:
+        """Promote a completed tuning run to create a deployable model.
+
+        Args:
+            run_id: Unique identifier of the tuning run to promote.
+
+        Returns:
+            The promotion result, including the created model ID.
+
+        Raises:
+            NotFoundError: If the tuning run does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        run_id = require_str(run_id, "run_id")
+        data = await self._request(
+            "POST", f"/v1/tuning-runs/{run_id}:promote", request_options=request_options
+        )
+        return self._maybe_validate(data, "PromoteResponse")
+
+    async def get_eval_run(
+        self,
+        eval_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> EvalRunDetailResponse:
+        """Retrieve details of an evaluation run.
+
+        Args:
+            eval_id: Unique identifier of the evaluation run.
+
+        Returns:
+            Detailed evaluation run information including metrics and
+            status.
+
+        Raises:
+            NotFoundError: If the evaluation run does not exist.
+            Knowledge2Error: If the API request fails.
+        """
+        eval_id = require_str(eval_id, "eval_id")
+        data = await self._request(
+            "GET", f"/v1/eval-runs/{eval_id}", request_options=request_options
+        )
+        return self._maybe_validate(data, "EvalRunDetailResponse")