# This file was auto-generated by Fern from our API Definition.

import typing
from json.decoder import JSONDecodeError

from ...core.api_error import ApiError
from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...core.http_response import AsyncHttpResponse, HttpResponse
from ...core.pydantic_utilities import parse_obj_as
from ...core.request_options import RequestOptions
from ...core.serialization import convert_and_respect_annotation_metadata
from ...errors.unprocessable_entity_error import UnprocessableEntityError
from ...types.chunk import Chunk
from ...types.prompt_message import PromptMessage
from ...types.structured_data_extractor_response import StructuredDataExtractorResponse

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)


class RawStructuredDataExtractorClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._client_wrapper = client_wrapper

    def invoke(
        self,
        *,
        chunks: typing.Sequence[Chunk],
        json_schema: typing.Dict[str, typing.Optional[typing.Any]],
        chunk_messages: typing.Optional[typing.Sequence[PromptMessage]] = OMIT,
        reduce: typing.Optional[bool] = OMIT,
        reduce_messages: typing.Optional[typing.Sequence[PromptMessage]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> HttpResponse[StructuredDataExtractorResponse]:
        """
        Extract structured data.

        tl;dr:
        - pass a valid JSON schema in `json_schema`
        - pass the page chunks as a list of `Chunk` objects, by default: `{"type": "text", "content": "..."}`
        - leave all other fields as default

        Detailed configuration (only relevant for complex use cases):

        The structured data extractor's architecture follows the map-reduce pattern,
        where the asset is divided into chunks, the schema is extracted from each chunk,
        and the chunks are then reduced to a single structured data object.

        In some applications, you may not want to:

        - map (if your input asset is small enough)
        - reduce (if your output object is large enough that it will overflow the output length;
            if you're extracting a long list of entities; if youre )
            to extract all instances of the schema).

        You can configure these behaviors with the `map` and `reduce` fields.

        Parameters
        ----------
        chunks : typing.Sequence[Chunk]
            The chunks from which to extract structured data.

        json_schema : typing.Dict[str, typing.Optional[typing.Any]]
            The JSON schema to use for validation (version draft 2020-12). See the docs [here](https://json-schema.org/learn/getting-started-step-by-step).

        chunk_messages : typing.Optional[typing.Sequence[PromptMessage]]
            The prompt to use for the data extraction over *each individual chunk*. It must be a list of messages.  The chunk content will be appended as a list of human messages.

        reduce : typing.Optional[bool]
            If `map`, whether to reduce the chunks to a single structured object (true) or return the full list (false).  Use True unless you want to preserve duplicates from each page or expect the object to overflow the output context.

        reduce_messages : typing.Optional[typing.Sequence[PromptMessage]]
            The prompt to use for the reduce steps. It must be a list of messages. The two extraction attempts will be appended as a list of human messages.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        HttpResponse[StructuredDataExtractorResponse]
            Successful Response
        """
        _response = self._client_wrapper.httpx_client.request(
            "api/v0/tools/structured-data-extractor/invoke",
            method="POST",
            json={
                "chunk_messages": convert_and_respect_annotation_metadata(
                    object_=chunk_messages, annotation=typing.Sequence[PromptMessage], direction="write"
                ),
                "chunks": convert_and_respect_annotation_metadata(
                    object_=chunks, annotation=typing.Sequence[Chunk], direction="write"
                ),
                "json_schema": json_schema,
                "reduce": reduce,
                "reduce_messages": convert_and_respect_annotation_metadata(
                    object_=reduce_messages, annotation=typing.Sequence[PromptMessage], direction="write"
                ),
            },
            headers={
                "content-type": "application/json",
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                _data = typing.cast(
                    StructuredDataExtractorResponse,
                    parse_obj_as(
                        type_=StructuredDataExtractorResponse,  # type: ignore
                        object_=_response.json(),
                    ),
                )
                return HttpResponse(response=_response, data=_data)
            if _response.status_code == 422:
                raise UnprocessableEntityError(
                    headers=dict(_response.headers),
                    body=typing.cast(
                        typing.Optional[typing.Any],
                        parse_obj_as(
                            type_=typing.Optional[typing.Any],  # type: ignore
                            object_=_response.json(),
                        ),
                    ),
                )
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)


class AsyncRawStructuredDataExtractorClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._client_wrapper = client_wrapper

    async def invoke(
        self,
        *,
        chunks: typing.Sequence[Chunk],
        json_schema: typing.Dict[str, typing.Optional[typing.Any]],
        chunk_messages: typing.Optional[typing.Sequence[PromptMessage]] = OMIT,
        reduce: typing.Optional[bool] = OMIT,
        reduce_messages: typing.Optional[typing.Sequence[PromptMessage]] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> AsyncHttpResponse[StructuredDataExtractorResponse]:
        """
        Extract structured data.

        tl;dr:
        - pass a valid JSON schema in `json_schema`
        - pass the page chunks as a list of `Chunk` objects, by default: `{"type": "text", "content": "..."}`
        - leave all other fields as default

        Detailed configuration (only relevant for complex use cases):

        The structured data extractor's architecture follows the map-reduce pattern,
        where the asset is divided into chunks, the schema is extracted from each chunk,
        and the chunks are then reduced to a single structured data object.

        In some applications, you may not want to:

        - map (if your input asset is small enough)
        - reduce (if your output object is large enough that it will overflow the output length;
            if you're extracting a long list of entities; if youre )
            to extract all instances of the schema).

        You can configure these behaviors with the `map` and `reduce` fields.

        Parameters
        ----------
        chunks : typing.Sequence[Chunk]
            The chunks from which to extract structured data.

        json_schema : typing.Dict[str, typing.Optional[typing.Any]]
            The JSON schema to use for validation (version draft 2020-12). See the docs [here](https://json-schema.org/learn/getting-started-step-by-step).

        chunk_messages : typing.Optional[typing.Sequence[PromptMessage]]
            The prompt to use for the data extraction over *each individual chunk*. It must be a list of messages.  The chunk content will be appended as a list of human messages.

        reduce : typing.Optional[bool]
            If `map`, whether to reduce the chunks to a single structured object (true) or return the full list (false).  Use True unless you want to preserve duplicates from each page or expect the object to overflow the output context.

        reduce_messages : typing.Optional[typing.Sequence[PromptMessage]]
            The prompt to use for the reduce steps. It must be a list of messages. The two extraction attempts will be appended as a list of human messages.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        AsyncHttpResponse[StructuredDataExtractorResponse]
            Successful Response
        """
        _response = await self._client_wrapper.httpx_client.request(
            "api/v0/tools/structured-data-extractor/invoke",
            method="POST",
            json={
                "chunk_messages": convert_and_respect_annotation_metadata(
                    object_=chunk_messages, annotation=typing.Sequence[PromptMessage], direction="write"
                ),
                "chunks": convert_and_respect_annotation_metadata(
                    object_=chunks, annotation=typing.Sequence[Chunk], direction="write"
                ),
                "json_schema": json_schema,
                "reduce": reduce,
                "reduce_messages": convert_and_respect_annotation_metadata(
                    object_=reduce_messages, annotation=typing.Sequence[PromptMessage], direction="write"
                ),
            },
            headers={
                "content-type": "application/json",
            },
            request_options=request_options,
            omit=OMIT,
        )
        try:
            if 200 <= _response.status_code < 300:
                _data = typing.cast(
                    StructuredDataExtractorResponse,
                    parse_obj_as(
                        type_=StructuredDataExtractorResponse,  # type: ignore
                        object_=_response.json(),
                    ),
                )
                return AsyncHttpResponse(response=_response, data=_data)
            if _response.status_code == 422:
                raise UnprocessableEntityError(
                    headers=dict(_response.headers),
                    body=typing.cast(
                        typing.Optional[typing.Any],
                        parse_obj_as(
                            type_=typing.Optional[typing.Any],  # type: ignore
                            object_=_response.json(),
                        ),
                    ),
                )
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response.text)
        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
