# ----------------------------------------------------------------------------
#
#  Welcome to Baml! To use this generated code, please run the following:
#
#  $ pip install baml
#
# ----------------------------------------------------------------------------

# This file was generated by BAML: please do not edit it. Instead, edit the
# BAML files and re-generate this code using: baml-cli generate
# baml-cli is available with the baml package.

import typing
import typing_extensions
import baml_py

from . import stream_types, types, type_builder
from .parser import LlmResponseParser, LlmStreamParser
from .runtime import DoNotUseDirectlyCallManager, BamlCallOptions
from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME as __runtime__


class BamlSyncClient:
    __options: DoNotUseDirectlyCallManager
    __stream_client: "BamlStreamClient"
    __http_request: "BamlHttpRequestClient"
    __http_stream_request: "BamlHttpStreamRequestClient"
    __llm_response_parser: LlmResponseParser
    __llm_stream_parser: LlmStreamParser

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options
        self.__stream_client = BamlStreamClient(options)
        self.__http_request = BamlHttpRequestClient(options)
        self.__http_stream_request = BamlHttpStreamRequestClient(options)
        self.__llm_response_parser = LlmResponseParser(options)
        self.__llm_stream_parser = LlmStreamParser(options)

    def __getstate__(self):
        # Return state needed for pickling
        return {"options": self.__options}

    def __setstate__(self, state):
        # Restore state from pickling
        self.__options = state["options"]
        self.__stream_client = BamlStreamClient(self.__options)
        self.__http_request = BamlHttpRequestClient(self.__options)
        self.__http_stream_request = BamlHttpStreamRequestClient(self.__options)
        self.__llm_response_parser = LlmResponseParser(self.__options)
        self.__llm_stream_parser = LlmStreamParser(self.__options)

    def with_options(
        self,
        tb: typing.Optional[type_builder.TypeBuilder] = None,
        client_registry: typing.Optional[baml_py.baml_py.ClientRegistry] = None,
        collector: typing.Optional[
            typing.Union[baml_py.baml_py.Collector, typing.List[baml_py.baml_py.Collector]]
        ] = None,
        env: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None,
    ) -> "BamlSyncClient":
        options: BamlCallOptions = {}
        if tb is not None:
            options["tb"] = tb
        if client_registry is not None:
            options["client_registry"] = client_registry
        if collector is not None:
            options["collector"] = collector
        if env is not None:
            options["env"] = env
        return BamlSyncClient(self.__options.merge_options(options))

    @property
    def stream(self):
        return self.__stream_client

    @property
    def request(self):
        return self.__http_request

    @property
    def stream_request(self):
        return self.__http_stream_request

    @property
    def parse(self):
        return self.__llm_response_parser

    @property
    def parse_stream(self):
        return self.__llm_stream_parser

    def ExtractCategories(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> types.DefaultContentPrediction:
        result = self.__options.merge_options(baml_options).call_function_sync(
            function_name="ExtractCategories",
            args={
                "content": content,
            },
        )
        return typing.cast(
            types.DefaultContentPrediction,
            result.cast_to(types, types, stream_types, False, __runtime__),
        )

    def ExtractContentGraphGeneric(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> types.KnowledgeGraph:
        result = self.__options.merge_options(baml_options).call_function_sync(
            function_name="ExtractContentGraphGeneric",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
        )
        return typing.cast(
            types.KnowledgeGraph, result.cast_to(types, types, stream_types, False, __runtime__)
        )

    def ExtractDynamicContentGraph(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> types.DynamicKnowledgeGraph:
        result = self.__options.merge_options(baml_options).call_function_sync(
            function_name="ExtractDynamicContentGraph",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
        )
        return typing.cast(
            types.DynamicKnowledgeGraph,
            result.cast_to(types, types, stream_types, False, __runtime__),
        )

    def SummarizeCode(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> types.SummarizedCode:
        result = self.__options.merge_options(baml_options).call_function_sync(
            function_name="SummarizeCode",
            args={
                "content": content,
            },
        )
        return typing.cast(
            types.SummarizedCode, result.cast_to(types, types, stream_types, False, __runtime__)
        )

    def SummarizeContent(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> types.SummarizedContent:
        result = self.__options.merge_options(baml_options).call_function_sync(
            function_name="SummarizeContent",
            args={
                "content": content,
            },
        )
        return typing.cast(
            types.SummarizedContent, result.cast_to(types, types, stream_types, False, __runtime__)
        )


class BamlStreamClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    def ExtractCategories(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlSyncStream[
        stream_types.DefaultContentPrediction, types.DefaultContentPrediction
    ]:
        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
            function_name="ExtractCategories",
            args={
                "content": content,
            },
        )
        return baml_py.BamlSyncStream[
            stream_types.DefaultContentPrediction, types.DefaultContentPrediction
        ](
            result,
            lambda x: typing.cast(
                stream_types.DefaultContentPrediction,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.DefaultContentPrediction,
                x.cast_to(types, types, stream_types, False, __runtime__),
            ),
            ctx,
        )

    def ExtractContentGraphGeneric(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph]:
        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
            function_name="ExtractContentGraphGeneric",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
        )
        return baml_py.BamlSyncStream[stream_types.KnowledgeGraph, types.KnowledgeGraph](
            result,
            lambda x: typing.cast(
                stream_types.KnowledgeGraph,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.KnowledgeGraph, x.cast_to(types, types, stream_types, False, __runtime__)
            ),
            ctx,
        )

    def ExtractDynamicContentGraph(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlSyncStream[stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph]:
        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
            function_name="ExtractDynamicContentGraph",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
        )
        return baml_py.BamlSyncStream[
            stream_types.DynamicKnowledgeGraph, types.DynamicKnowledgeGraph
        ](
            result,
            lambda x: typing.cast(
                stream_types.DynamicKnowledgeGraph,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.DynamicKnowledgeGraph,
                x.cast_to(types, types, stream_types, False, __runtime__),
            ),
            ctx,
        )

    def SummarizeCode(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode]:
        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
            function_name="SummarizeCode",
            args={
                "content": content,
            },
        )
        return baml_py.BamlSyncStream[stream_types.SummarizedCode, types.SummarizedCode](
            result,
            lambda x: typing.cast(
                stream_types.SummarizedCode,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.SummarizedCode, x.cast_to(types, types, stream_types, False, __runtime__)
            ),
            ctx,
        )

    def SummarizeContent(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent]:
        ctx, result = self.__options.merge_options(baml_options).create_sync_stream(
            function_name="SummarizeContent",
            args={
                "content": content,
            },
        )
        return baml_py.BamlSyncStream[stream_types.SummarizedContent, types.SummarizedContent](
            result,
            lambda x: typing.cast(
                stream_types.SummarizedContent,
                x.cast_to(types, types, stream_types, True, __runtime__),
            ),
            lambda x: typing.cast(
                types.SummarizedContent, x.cast_to(types, types, stream_types, False, __runtime__)
            ),
            ctx,
        )


class BamlHttpRequestClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    def ExtractCategories(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractCategories",
            args={
                "content": content,
            },
            mode="request",
        )
        return result

    def ExtractContentGraphGeneric(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractContentGraphGeneric",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
            mode="request",
        )
        return result

    def ExtractDynamicContentGraph(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractDynamicContentGraph",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
            mode="request",
        )
        return result

    def SummarizeCode(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="SummarizeCode",
            args={
                "content": content,
            },
            mode="request",
        )
        return result

    def SummarizeContent(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="SummarizeContent",
            args={
                "content": content,
            },
            mode="request",
        )
        return result


class BamlHttpStreamRequestClient:
    __options: DoNotUseDirectlyCallManager

    def __init__(self, options: DoNotUseDirectlyCallManager):
        self.__options = options

    def ExtractCategories(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractCategories",
            args={
                "content": content,
            },
            mode="stream",
        )
        return result

    def ExtractContentGraphGeneric(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractContentGraphGeneric",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
            mode="stream",
        )
        return result

    def ExtractDynamicContentGraph(
        self,
        content: str,
        mode: typing.Optional[
            typing.Union[
                typing_extensions.Literal["simple"],
                typing_extensions.Literal["base"],
                typing_extensions.Literal["guided"],
                typing_extensions.Literal["strict"],
                typing_extensions.Literal["custom"],
            ]
        ] = None,
        custom_prompt_content: typing.Optional[str] = None,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="ExtractDynamicContentGraph",
            args={
                "content": content,
                "mode": mode,
                "custom_prompt_content": custom_prompt_content,
            },
            mode="stream",
        )
        return result

    def SummarizeCode(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="SummarizeCode",
            args={
                "content": content,
            },
            mode="stream",
        )
        return result

    def SummarizeContent(
        self,
        content: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = self.__options.merge_options(baml_options).create_http_request_sync(
            function_name="SummarizeContent",
            args={
                "content": content,
            },
            mode="stream",
        )
        return result


b = BamlSyncClient(DoNotUseDirectlyCallManager({}))
