# ----------------------------------------------------------------------------
#
#  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 BamlAsyncClient:
    __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 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,
    ) -> "BamlAsyncClient":
        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 BamlAsyncClient(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
    
    async def AbstractPage(self, page: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> types.PageAbstract:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="AbstractPage", args={
            "page": page,
        })
        return typing.cast(types.PageAbstract, result.cast_to(types, types, stream_types, False, __runtime__))
    async def ExtractFromElement(self, screenshot: baml_py.Image,html: str,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> types.Output:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractFromElement", args={
            "screenshot": screenshot,"html": html,"instruction": instruction,
        })
        return typing.cast(types.Output, result.cast_to(types, types, stream_types, False, __runtime__))
    async def ExtractFromState(self, screenshot: baml_py.Image,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> types.Output:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractFromState", args={
            "screenshot": screenshot,"instruction": instruction,
        })
        return typing.cast(types.Output, result.cast_to(types, types, stream_types, False, __runtime__))
    async def ExtractTestCase(self, description: str,
        baml_options: BamlCallOptions = {},
    ) -> types.TestCase:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ExtractTestCase", args={
            "description": description,
        })
        return typing.cast(types.TestCase, result.cast_to(types, types, stream_types, False, __runtime__))
    async def GeneratePostcondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> str:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="GeneratePostcondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        })
        return typing.cast(str, result.cast_to(types, types, stream_types, False, __runtime__))
    async def GeneratePrecondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> str:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="GeneratePrecondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        })
        return typing.cast(str, result.cast_to(types, types, stream_types, False, __runtime__))
    async def ImplicitGeneration(self, testcase: types.TestCase,
        baml_options: BamlCallOptions = {},
    ) -> types.TestCase:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ImplicitGeneration", args={
            "testcase": testcase,
        })
        return typing.cast(types.TestCase, result.cast_to(types, types, stream_types, False, __runtime__))
    async def IsSameLogicalPage(self, page_a: baml_py.Image,page_b: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> bool:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="IsSameLogicalPage", args={
            "page_a": page_a,"page_b": page_b,
        })
        return typing.cast(bool, result.cast_to(types, types, stream_types, False, __runtime__))
    async def LocateUIElement(self, screenshot: baml_py.Image,description: str,
        baml_options: BamlCallOptions = {},
    ) -> str:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="LocateUIElement", args={
            "screenshot": screenshot,"description": description,
        })
        return typing.cast(str, result.cast_to(types, types, stream_types, False, __runtime__))
    async def ProposeActions(self, screenshot: baml_py.Image,task: str,
        baml_options: BamlCallOptions = {},
    ) -> str:
        result = await self.__options.merge_options(baml_options).call_function_async(function_name="ProposeActions", args={
            "screenshot": screenshot,"task": task,
        })
        return typing.cast(str, result.cast_to(types, types, stream_types, False, __runtime__))
    


class BamlStreamClient:
    __options: DoNotUseDirectlyCallManager

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

    def AbstractPage(self, page: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.PageAbstract, types.PageAbstract]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="AbstractPage", args={
            "page": page,
        })
        return baml_py.BamlStream[stream_types.PageAbstract, types.PageAbstract](
          result,
          lambda x: typing.cast(stream_types.PageAbstract, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(types.PageAbstract, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def ExtractFromElement(self, screenshot: baml_py.Image,html: str,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.Output, types.Output]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractFromElement", args={
            "screenshot": screenshot,"html": html,"instruction": instruction,
        })
        return baml_py.BamlStream[stream_types.Output, types.Output](
          result,
          lambda x: typing.cast(stream_types.Output, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(types.Output, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def ExtractFromState(self, screenshot: baml_py.Image,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.Output, types.Output]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractFromState", args={
            "screenshot": screenshot,"instruction": instruction,
        })
        return baml_py.BamlStream[stream_types.Output, types.Output](
          result,
          lambda x: typing.cast(stream_types.Output, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(types.Output, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def ExtractTestCase(self, description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.TestCase, types.TestCase]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ExtractTestCase", args={
            "description": description,
        })
        return baml_py.BamlStream[stream_types.TestCase, types.TestCase](
          result,
          lambda x: typing.cast(stream_types.TestCase, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(types.TestCase, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def GeneratePostcondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[str, str]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="GeneratePostcondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        })
        return baml_py.BamlStream[str, str](
          result,
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def GeneratePrecondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[str, str]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="GeneratePrecondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        })
        return baml_py.BamlStream[str, str](
          result,
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def ImplicitGeneration(self, testcase: types.TestCase,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[stream_types.TestCase, types.TestCase]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ImplicitGeneration", args={
            "testcase": testcase,
        })
        return baml_py.BamlStream[stream_types.TestCase, types.TestCase](
          result,
          lambda x: typing.cast(stream_types.TestCase, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(types.TestCase, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def IsSameLogicalPage(self, page_a: baml_py.Image,page_b: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[bool, bool]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="IsSameLogicalPage", args={
            "page_a": page_a,"page_b": page_b,
        })
        return baml_py.BamlStream[bool, bool](
          result,
          lambda x: typing.cast(bool, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(bool, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def LocateUIElement(self, screenshot: baml_py.Image,description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[str, str]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="LocateUIElement", args={
            "screenshot": screenshot,"description": description,
        })
        return baml_py.BamlStream[str, str](
          result,
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    def ProposeActions(self, screenshot: baml_py.Image,task: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.BamlStream[str, str]:
        ctx, result = self.__options.merge_options(baml_options).create_async_stream(function_name="ProposeActions", args={
            "screenshot": screenshot,"task": task,
        })
        return baml_py.BamlStream[str, str](
          result,
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, True, __runtime__)),
          lambda x: typing.cast(str, x.cast_to(types, types, stream_types, False, __runtime__)),
          ctx,
        )
    

class BamlHttpRequestClient:
    __options: DoNotUseDirectlyCallManager

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

    async def AbstractPage(self, page: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="AbstractPage", args={
            "page": page,
        }, mode="request")
        return result
    async def ExtractFromElement(self, screenshot: baml_py.Image,html: str,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractFromElement", args={
            "screenshot": screenshot,"html": html,"instruction": instruction,
        }, mode="request")
        return result
    async def ExtractFromState(self, screenshot: baml_py.Image,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractFromState", args={
            "screenshot": screenshot,"instruction": instruction,
        }, mode="request")
        return result
    async def ExtractTestCase(self, description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractTestCase", args={
            "description": description,
        }, mode="request")
        return result
    async def GeneratePostcondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="GeneratePostcondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        }, mode="request")
        return result
    async def GeneratePrecondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="GeneratePrecondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        }, mode="request")
        return result
    async def ImplicitGeneration(self, testcase: types.TestCase,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ImplicitGeneration", args={
            "testcase": testcase,
        }, mode="request")
        return result
    async def IsSameLogicalPage(self, page_a: baml_py.Image,page_b: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="IsSameLogicalPage", args={
            "page_a": page_a,"page_b": page_b,
        }, mode="request")
        return result
    async def LocateUIElement(self, screenshot: baml_py.Image,description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="LocateUIElement", args={
            "screenshot": screenshot,"description": description,
        }, mode="request")
        return result
    async def ProposeActions(self, screenshot: baml_py.Image,task: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ProposeActions", args={
            "screenshot": screenshot,"task": task,
        }, mode="request")
        return result
    

class BamlHttpStreamRequestClient:
    __options: DoNotUseDirectlyCallManager

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

    async def AbstractPage(self, page: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="AbstractPage", args={
            "page": page,
        }, mode="stream")
        return result
    async def ExtractFromElement(self, screenshot: baml_py.Image,html: str,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractFromElement", args={
            "screenshot": screenshot,"html": html,"instruction": instruction,
        }, mode="stream")
        return result
    async def ExtractFromState(self, screenshot: baml_py.Image,instruction: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractFromState", args={
            "screenshot": screenshot,"instruction": instruction,
        }, mode="stream")
        return result
    async def ExtractTestCase(self, description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ExtractTestCase", args={
            "description": description,
        }, mode="stream")
        return result
    async def GeneratePostcondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="GeneratePostcondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        }, mode="stream")
        return result
    async def GeneratePrecondition(self, screenshot: baml_py.Image,history: typing.List["types.History"],action: str,verify: str,feedback: typing.List["types.Feedback"],
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="GeneratePrecondition", args={
            "screenshot": screenshot,"history": history,"action": action,"verify": verify,"feedback": feedback,
        }, mode="stream")
        return result
    async def ImplicitGeneration(self, testcase: types.TestCase,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ImplicitGeneration", args={
            "testcase": testcase,
        }, mode="stream")
        return result
    async def IsSameLogicalPage(self, page_a: baml_py.Image,page_b: baml_py.Image,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="IsSameLogicalPage", args={
            "page_a": page_a,"page_b": page_b,
        }, mode="stream")
        return result
    async def LocateUIElement(self, screenshot: baml_py.Image,description: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="LocateUIElement", args={
            "screenshot": screenshot,"description": description,
        }, mode="stream")
        return result
    async def ProposeActions(self, screenshot: baml_py.Image,task: str,
        baml_options: BamlCallOptions = {},
    ) -> baml_py.baml_py.HTTPRequest:
        result = await self.__options.merge_options(baml_options).create_http_request_async(function_name="ProposeActions", args={
            "screenshot": screenshot,"task": task,
        }, mode="stream")
        return result
    

b = BamlAsyncClient(DoNotUseDirectlyCallManager({}))