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

import typing

from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...core.request_options import RequestOptions
from ...types.chunk import Chunk
from ...types.prompt_message import PromptMessage
from ...types.structured_data_extractor_response import StructuredDataExtractorResponse
from .raw_client import AsyncRawStructuredDataExtractorClient, RawStructuredDataExtractorClient

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


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

    @property
    def with_raw_response(self) -> RawStructuredDataExtractorClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        RawStructuredDataExtractorClient
        """
        return self._raw_client

    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,
    ) -> 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
        -------
        StructuredDataExtractorResponse
            Successful Response

        Examples
        --------
        from athena import Athena, Chunk, ChunkContentItem_Text

        client = Athena(
            api_key="YOUR_API_KEY",
        )
        client.tools.structured_data_extractor.invoke(
            chunks=[
                Chunk(
                    chunk_id="1",
                    content=[
                        ChunkContentItem_Text(
                            text="John Smith is a 35 year old developer. You can reach him at john.smith@example.com",
                        )
                    ],
                ),
                Chunk(
                    chunk_id="2",
                    content=[
                        ChunkContentItem_Text(
                            text="Jane Doe is a 25 year old developer. You can reach her at jane@example.com",
                        )
                    ],
                ),
            ],
            json_schema={
                "description": "A person",
                "properties": {
                    "age": {"type": "integer"},
                    "email": {"type": "string"},
                    "name": {"type": "string"},
                },
                "required": ["name"],
                "title": "Person",
                "type": "object",
            },
        )
        """
        _response = self._raw_client.invoke(
            chunks=chunks,
            json_schema=json_schema,
            chunk_messages=chunk_messages,
            reduce=reduce,
            reduce_messages=reduce_messages,
            request_options=request_options,
        )
        return _response.data


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

    @property
    def with_raw_response(self) -> AsyncRawStructuredDataExtractorClient:
        """
        Retrieves a raw implementation of this client that returns raw responses.

        Returns
        -------
        AsyncRawStructuredDataExtractorClient
        """
        return self._raw_client

    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,
    ) -> 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
        -------
        StructuredDataExtractorResponse
            Successful Response

        Examples
        --------
        import asyncio

        from athena import AsyncAthena, Chunk, ChunkContentItem_Text

        client = AsyncAthena(
            api_key="YOUR_API_KEY",
        )


        async def main() -> None:
            await client.tools.structured_data_extractor.invoke(
                chunks=[
                    Chunk(
                        chunk_id="1",
                        content=[
                            ChunkContentItem_Text(
                                text="John Smith is a 35 year old developer. You can reach him at john.smith@example.com",
                            )
                        ],
                    ),
                    Chunk(
                        chunk_id="2",
                        content=[
                            ChunkContentItem_Text(
                                text="Jane Doe is a 25 year old developer. You can reach her at jane@example.com",
                            )
                        ],
                    ),
                ],
                json_schema={
                    "description": "A person",
                    "properties": {
                        "age": {"type": "integer"},
                        "email": {"type": "string"},
                        "name": {"type": "string"},
                    },
                    "required": ["name"],
                    "title": "Person",
                    "type": "object",
                },
            )


        asyncio.run(main())
        """
        _response = await self._raw_client.invoke(
            chunks=chunks,
            json_schema=json_schema,
            chunk_messages=chunk_messages,
            reduce=reduce,
            reduce_messages=reduce_messages,
            request_options=request_options,
        )
        return _response.data
