# This file is automatically generated by pyo3_stub_gen and deepseek-R1
# ruff: noqa: E501, F401, F811

import builtins
import typing

try:
    import numpy
    import numpy.typing as npt

    _NDArrayF32 = npt.NDArray[numpy.float32]
    _NDArrayAny = npt.NDArray[
        typing.Any
    ]  # For a more general ndarray if dtype is not float32
except ImportError:
    _NDArrayF32 = (
        typing.Any
    )  # Fallback if numpy is not available for type checking environment
    _NDArrayAny = typing.Any

class OpenAIEmbeddingData:
    """
    Represents a single embedding object.

    Attributes:
        object: The object type (usually a string identifier).
        index: Position of the embedding in the response.
    """

    object: builtins.str
    index: builtins.int

    @property
    def embedding(self) -> typing.Union[builtins.str, builtins.list[builtins.float]]:
        """
        The embedding vector. It may be either a list of floats or a base64 encoded string.

        Returns:
            Either a list of floats (when the embedding is decoded) or a string.
        """
        ...

class OpenAIUsage:
    """
    Represents token usage information for an API request.

    Attributes:
        prompt_tokens: The number of tokens used for the prompt.
        total_tokens: The total number of tokens used.
    """

    prompt_tokens: builtins.int
    total_tokens: builtins.int

class OpenAIEmbeddingsResponse:
    """
    Represents the response returned by an embeddings API request.

    Attributes:
        object: The object type (e.g. "list").
        data: A list of OpenAIEmbeddingData objects.
        model: The model identifier used for the embedding.
        usage: Usage details such as token counts.

    Methods:
        numpy() -> _NDArrayF32:
            Converts the embeddings data into a 2D NumPy array.

    Example:
        >>> response = client.embed(["Hello"], model="model-id")
        >>> array = response.numpy()
    """

    object: builtins.str
    data: builtins.list[OpenAIEmbeddingData]
    model: builtins.str
    usage: OpenAIUsage

    def numpy(self) -> _NDArrayF32:
        """
        Converts the embeddings data into a 2D NumPy array.

        Each row in the array corresponds to an embedding. The array's type will
        be float32 and its shape will be (number_of_embeddings, embedding_dimension).

        Returns:
            numpy.ndarray: A 2D NumPy array of f32.

        Raises:
            ValueError: If an embedding is not a float vector, if there are inconsistent dimensions,
                        or if the data is empty.
            ImportError: If NumPy is not installed in the Python environment.

        Example:
            >>> array = response.numpy()
            >>> print(array.shape)
        """
        ...

class RerankResult:
    """
    Represents a single result from a rerank operation.

    Attributes:
        index: The index of the result.
        score: The rerank score.
        text: The optional original text.
    """

    index: builtins.int
    score: builtins.float
    text: typing.Optional[builtins.str]

class RerankResponse:
    """
    Represents the response for a rerank API request.

    Attributes:
        object: The object type (usually "list").
        data: A list of RerankResult objects.

    Example:
        >>> response = client.rerank("query", ["doc1", "doc2"])
        >>> for result in response.data:
        ...     print(result.index, result.score)
    """

    object: builtins.str
    data: builtins.list[RerankResult]

    def __init__(self, data: builtins.list[RerankResult]) -> None:
        """
        Initializes a RerankResponse.

        Args:
            data: A list of RerankResult objects.
        """
        ...

class ClassificationResult:
    """
    Represents a single classification result.

    Attributes:
        label: The predicted classification label.
        score: The confidence score.
    """

    label: builtins.str
    score: builtins.float

class ClassificationResponse:
    """
    Represents the response from a classification API request.

    Attributes:
        object: The object type (usually "list").
        data: A nested list of ClassificationResult objects.

    Example:
        >>> response = client.classify(["text1", "text2"])
        >>> for result_group in response.data:
        ...     for result in result_group:
        ...         print(result.label, result.score)
    """

    object: builtins.str
    data: builtins.list[builtins.list[ClassificationResult]]

    def __init__(
        self, data: builtins.list[builtins.list[ClassificationResult]]
    ) -> None:
        """
        Initializes a ClassificationResponse.

        Args:
            data: A list where each element is a list of ClassificationResult objects.
        """
        ...

class InferenceClient:
    """
    Baseten.co API client for embedding, reranking, and classification.

    This client allows you to send text to an embedding model, rerank documents,
    or classify texts through the API.

    Attributes:
        api_base: The base URL for the API.
        api_key: The API key for authentication.

    Example:
        >>> client = InferenceClient(api_base="https://example.api.baseten.co/environments/production/sync", api_key="your_api_key")
        >>> embeddings = client.embed(["Hello world"], model="BAAI/bge-large-en")
        >>> array = embeddings.numpy()
    """
    def __init__(
        self, api_base: builtins.str, api_key: typing.Optional[builtins.str] = None
    ) -> None:
        """
        Initialize the sync client with the API base URL and optional API key.

        Args:
            api_base: The base URL of the API.
            api_key: The API key. If not provided, environment variables will be checked.

        Example:
            >>> client = InferenceClient(api_base="https://example.api.baseten.co/sync", api_key="your_key")
        """
        ...

    @property
    def api_key(self) -> builtins.str:
        """
        Returns the configured API key.

        Example:
            >>> print(client.api_key)
        """
        ...

    def embed(
        self,
        input: builtins.list[builtins.str],
        model: builtins.str,
        encoding_format: typing.Optional[builtins.str] = None,
        dimensions: typing.Optional[builtins.int] = None,
        user: typing.Optional[builtins.str] = None,
        max_concurrent_requests: builtins.int = 32,  # DEFAULT_CONCURRENCY
        batch_size: builtins.int = 16,  # DEFAULT_BATCH_SIZE
        timeout_s: builtins.float = 3600.0,  # DEFAULT_REQUEST_TIMEOUT_S
    ) -> OpenAIEmbeddingsResponse:
        """
        Sends a list of strings to the embedding endpoint to generate embeddings.

        Args:
            input: A list of texts to embed.
            model: The model identifier.
            encoding_format: Optional encoding format.
            dimensions: Optional dimension size of the embeddings.
            user: Optional user identifier.
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Number of texts per batch.
            timeout_s: Total timeout in seconds.

        Returns:
            An OpenAIEmbeddingsResponse object.

        Raises:
            ValueError: If the input list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = client.embed(["hello", "world"], model="model-id")
            >>> print(response.data[0].embedding)

        Example with error handling:
            >>> try:
            ...     response = client.embed(["hello", "world"], model="invalid-model")
            ... except requests.exceptions.HTTPError as e:
            ...     print(f"Request failed: {e} with status code {e.response.args[0]}")
        """
        ...

    def rerank(
        self,
        query: builtins.str,
        texts: builtins.list[builtins.str],
        raw_scores: builtins.bool = False,
        return_text: builtins.bool = False,
        truncate: builtins.bool = False,
        truncation_direction: builtins.str = "Right",
        max_concurrent_requests: builtins.int = 32,  # DEFAULT_CONCURRENCY
        batch_size: builtins.int = 16,  # DEFAULT_BATCH_SIZE
        timeout_s: builtins.float = 3600.0,  # DEFAULT_REQUEST_TIMEOUT_S
    ) -> RerankResponse:
        """
        Reranks a set of texts based on the provided query.

        Args:
            query: The query string.
            texts: A list of texts to rerank.
            raw_scores: Whether raw scores should be returned.
            return_text: Whether to include the original text.
            truncate: Whether to truncate texts.
            truncation_direction: Direction for truncation ('Right' by default).
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Batch size for each request.
            timeout_s: Overall timeout in seconds.

        Returns:
            A RerankResponse object.

        Raises:
            ValueError: If the texts list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = client.rerank("find", ["doc1", "doc2"])
            >>> for result in response.data:
            ...     print(result.index, result.score)
        """
        ...

    def classify(
        self,
        inputs: builtins.list[builtins.str],
        raw_scores: builtins.bool = False,
        truncate: builtins.bool = False,
        truncation_direction: builtins.str = "Right",
        max_concurrent_requests: builtins.int = 32,  # DEFAULT_CONCURRENCY
        batch_size: builtins.int = 16,  # DEFAULT_BATCH_SIZE
        timeout_s: builtins.float = 3600.0,  # DEFAULT_REQUEST_TIMEOUT_S
    ) -> ClassificationResponse:
        """
        Classifies each input text.

        Args:
            inputs: A list of texts to classify.
            raw_scores: Whether raw scores should be returned.
            truncate: Whether to truncate long texts.
            truncation_direction: Truncation direction ('Right' by default).
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Batch size for each request.
            timeout_s: Overall timeout in seconds.

        Returns:
            A ClassificationResponse object.

        Raises:
            ValueError: If the inputs list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = client.classify(["text1", "text2"])
            >>> for group in response.data:
            ...     for result in group:
            ...         print(result.label, result.score)
        """
        ...

    def batch_post(
        self,
        url_path: builtins.str,
        payloads: builtins.list[typing.Any],
        max_concurrent_requests: builtins.int = 32,  # DEFAULT_CONCURRENCY
        timeout_s: builtins.float = 3600.0,  # DEFAULT_REQUEST_TIMEOUT_S
    ) -> builtins.list[typing.Any]:
        """
        Sends a list of generic JSON payloads to a specified URL path concurrently.

        Each payload is sent as an individual POST request. The responses are
        returned as a list of Python objects, corresponding to the JSON responses
        from the server.

        Args:
            url_path: The specific API path to post to (e.g., "/v1/custom_endpoint").
            payloads: A list of Python objects that are JSON-serializable.
                      Each object will be the body of a POST request.
            max_concurrent_requests: Maximum number of parallel requests.
            timeout_s: Total timeout in seconds for the entire batch operation,
                       also used as the timeout for each individual request.

        Returns:
            A list of Python objects, where each object is the deserialized
            JSON response from the server for the corresponding request payload.
            The order of responses matches the order of input payloads.

        Raises:
            ValueError: If the payloads list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If any of the underlying HTTP requests fail.
            # Note: Other PyO3/Rust errors might be raised for serialization/deserialization issues.

        Example:
            >>> client = InferenceClient(api_base="https://example.api.baseten.co/sync", api_key="your_key")
            >>> custom_payloads = [
            ...     {"data": "request1_data", "id": 1},
            ...     {"data": "request2_data", "id": 2}
            ... ]
            >>> responses = client.batch_post("/v1/process_item", custom_payloads)
            >>> for resp in responses:
            ...     print(resp)
        """
        ...

    async def aembed(
        self,
        input: builtins.list[builtins.str],
        model: builtins.str,
        encoding_format: typing.Optional[builtins.str] = None,
        dimensions: typing.Optional[builtins.int] = None,
        user: typing.Optional[builtins.str] = None,
        max_concurrent_requests: builtins.int = 32,
        batch_size: builtins.int = 16,
        timeout_s: builtins.float = 3600.0,
    ) -> OpenAIEmbeddingsResponse:
        """
        Asynchronously sends a list of texts to the embedding endpoint to generate embeddings.

        Args:
            input: A list of texts to embed.
            model: The model identifier.
            encoding_format: Optional encoding format.
            dimensions: Optional dimension size of the embeddings.
            user: Optional user identifier.
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Number of texts per batch.
            timeout_s: Total timeout in seconds.

        Returns:
            An awaitable OpenAIEmbeddingsResponse object.

        Raises:
            ValueError: If the input list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = await client.aembed(["hello", "world"], model="model-id")
            >>> print(response.data[0].embedding)
        """
        ...

    async def arerank(
        self,
        query: builtins.str,
        texts: builtins.list[builtins.str],
        raw_scores: builtins.bool = False,
        return_text: builtins.bool = False,
        truncate: builtins.bool = False,
        truncation_direction: builtins.str = "Right",
        max_concurrent_requests: builtins.int = 32,
        batch_size: builtins.int = 16,
        timeout_s: builtins.float = 3600.0,
    ) -> RerankResponse:
        """
        Asynchronously reranks a set of texts based on the provided query.

        Args:
            query: The query string.
            texts: A list of texts to rerank.
            raw_scores: Whether raw scores should be returned.
            return_text: Whether to include the original text.
            truncate: Whether to truncate texts.
            truncation_direction: Direction for truncation ('Right' by default).
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Batch size for each request.
            timeout_s: Overall timeout in seconds.

        Returns:
            An awaitable RerankResponse object.

        Raises:
            ValueError: If the texts list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = await client.arerank("find", ["doc1", "doc2"])
            >>> for result in response.data:
            ...     print(result.index, result.score)
        """
        ...

    async def aclassify(
        self,
        inputs: builtins.list[builtins.str],
        raw_scores: builtins.bool = False,
        truncate: builtins.bool = False,
        truncation_direction: builtins.str = "Right",
        max_concurrent_requests: builtins.int = 32,
        batch_size: builtins.int = 16,
        timeout_s: builtins.float = 3600.0,
    ) -> ClassificationResponse:
        """
        Asynchronously classifies each input text.

        Args:
            inputs: A list of texts to classify.
            raw_scores: Whether raw scores should be returned.
            truncate: Whether to truncate long texts.
            truncation_direction: Truncation direction ('Right' by default).
            max_concurrent_requests: Maximum parallel requests.
            batch_size: Batch size for each request.
            timeout_s: Overall timeout in seconds.

        Returns:
            An awaitable ClassificationResponse object.

        Raises:
            ValueError: If the inputs list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If the request fails.

        Example:
            >>> response = await client.aclassify(["text1", "text2"])
            >>> for group in response.data:
            ...     for result in group:
            ...         print(result.label, result.score)
        """
        ...

    async def abatch_post(
        self,
        url_path: builtins.str,
        payloads: builtins.list[typing.Any],
        max_concurrent_requests: builtins.int = 32,
        timeout_s: builtins.float = 3600.0,
    ) -> builtins.list[typing.Any]:
        """
        Asynchronously sends a list of generic JSON payloads to a specified URL path concurrently.

        Args:
            url_path: The specific API path to post to (e.g., "/v1/custom_endpoint").
            payloads: A list of Python objects that are JSON-serializable.
            max_concurrent_requests: Maximum number of parallel requests.
            timeout_s: Total timeout in seconds for the batch operation.

        Returns:
            An awaitable list of Python objects whose order matches the input payloads.

        Raises:
            ValueError: If the payloads list is empty or parameters are invalid.
            requests.exceptions.HTTPError: If any underlying HTTP requests fail.

        Example:
            >>> responses = await client.abatch_post("/v1/process_item", [{"data": "r1"}, {"data": "r2"}])
            >>> for resp in responses:
            ...     print(resp)
        """
        ...

__version__: builtins.str
"""The version of the  bei_client library."""
