prompt_engine.py

from typing import Any, Dict, Generator, List, Literal, Union, cast

from typing_extensions import NotRequired, TypedDict

from ._config import ClientConfig
from .async_request import AsyncRequest
from .helpers import build_path
from .request import Request, RequestConfig


class PromptEngineResult(TypedDict):
    prompt: str
    return_prompt: str
    id: str


class PromptEngineRunParams(TypedDict):
    prompt: str
    inputs: NotRequired[List[object]]
    return_prompt: Union[str, List[object], Dict[str, str]]
    input_values: NotRequired[Dict[str, str]]
    stream: Union[bool, None] = False
    use_internet: Union[bool, None] = False
    prompt_guard: NotRequired[
        Literal[
            "defamation",
            "specialized_advice",
            "privacy",
            "intellectual_property",
            "indiscriminate_weapons",
            "hate",
            "sexual_content",
            "elections",
            "code_interpreter_abuse",
        ]
    ]


class PromptEngineExecuteParams(TypedDict):
    id: str
    input_values: object
    stream: Union[bool, None] = False


class PromptEngineRunResponse(TypedDict):
    success: bool
    result: Any


class PromptEngineCreateParams(TypedDict):
    name: NotRequired[str]
    prompt: str
    inputs: NotRequired[List[object]]
    return_prompt: Union[str, List[object], Dict[str, str]]
    use_internet: Union[bool, None] = False
    optimize_prompt: Union[bool, None] = False
    prompt_guard: NotRequired[
        Literal[
            "defamation",
            "specialized_advice",
            "privacy",
            "intellectual_property",
            "indiscriminate_weapons",
            "hate",
            "sexual_content",
            "elections",
            "code_interpreter_abuse",
        ]
    ]


class PromptEngineCreateResponse(TypedDict):
    success: bool
    prompt_engine_id: str


class PromptEngineGetResponse(PromptEngineResult):
    success: bool


class PromptEngineListResponse(TypedDict):
    success: bool
    prompt_engines: List[PromptEngineResult]


class PromptEngineListParams(TypedDict):
    limit: str
    page: str


class PromptEngineDeleteResponse(TypedDict):
    prompt_engine_id: str


class PromptEngine(ClientConfig):
    config: RequestConfig

    def __init__(
        self,
        api_key: str,
        base_url: str,
        headers: Union[Dict[str, str], None] = None,
    ):
        super().__init__(api_key, base_url, headers)
        self.config = RequestConfig(
            base_url=base_url,
            api_key=api_key,
            headers=headers,
        )

    def create(self, params: PromptEngineCreateParams) -> PromptEngineCreateResponse:
        path = "/prompt_engine"
        resp = Request(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp

    def get(self, id: str) -> PromptEngineGetResponse:
        path = f"/prompt_engine/{id}"
        resp = Request(config=self.config, path=path, params={}, verb="get").perform_with_content()
        return resp

    def list(self, params: Union[PromptEngineListParams, None] = None) -> PromptEngineListResponse:
        if params is None:
            params = {}

        # Default limit and page to 20 and 1 respectively
        if params.get("limit") is None:
            params["limit"] = 20

        if params.get("page") is None:
            params["page"] = 0

        path = build_path(
            base_path="/prompt_engine",
            params=params,
        )
        resp = Request(config=self.config, path=path, params={}, verb="get").perform_with_content()
        return resp

    def delete(self, id: str) -> PromptEngineDeleteResponse:
        path = f"/prompt_engine/{id}"
        resp = Request(
            config=self.config,
            path=path,
            params={},
            verb="DELETE",
        ).perform_with_content()
        return resp

    def run_prompt_direct(
        self, params: PromptEngineRunParams
    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
        path = "/prompt_engine/run"
        stream = params.get("stream")
        if stream:
            resp = Request(
                config=self.config,
                path=path,
                params=cast(Dict[Any, Any], params),
                verb="post",
            ).perform_with_content_streaming()
            return resp

        resp = Request(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp

    def run(
        self, params: PromptEngineExecuteParams
    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
        id = params.get("id")
        path = f"/prompt_engine/{id}"
        stream = params.get("stream")

        if stream:
            resp = Request(
                config=self.config,
                path=path,
                params=cast(Dict[Any, Any], params),
                verb="post",
            ).perform_with_content_streaming()
            return resp

        resp = Request(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp


class AsyncPromptEngine(ClientConfig):
    config: RequestConfig

    def __init__(
        self,
        api_key: str,
        base_url: str,
        headers: Union[Dict[str, str], None] = None,
    ):
        super().__init__(api_key, base_url, headers)
        self.config = RequestConfig(
            base_url=base_url,
            api_key=api_key,
            headers=headers,
        )

    async def create(self, params: PromptEngineCreateParams) -> PromptEngineCreateResponse:
        path = "/prompt_engine"
        resp = await AsyncRequest(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp

    async def get(self, id: str) -> PromptEngineGetResponse:
        path = f"/prompt_engine/{id}"
        resp = await AsyncRequest(
            config=self.config, path=path, params={}, verb="get"
        ).perform_with_content()
        return resp

    async def list(
        self, params: Union[PromptEngineListParams, None] = None
    ) -> PromptEngineListResponse:
        if params is None:
            params = {}

        # Default limit and page to 20 and 1 respectively
        if params.get("limit") is None:
            params["limit"] = 20

        if params.get("page") is None:
            params["page"] = 0

        path = build_path(
            base_path="/prompt_engine",
            params=params,
        )
        resp = await AsyncRequest(
            config=self.config, path=path, params={}, verb="get"
        ).perform_with_content()
        return resp

    async def delete(self, id: str) -> PromptEngineDeleteResponse:
        path = f"/prompt_engine/{id}"
        resp = await AsyncRequest(
            config=self.config,
            path=path,
            params={},
            verb="DELETE",
        ).perform_with_content()
        return resp

    async def run_prompt_direct(
        self, params: PromptEngineRunParams
    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
        path = "/prompt_engine/run"
        stream = params.get("stream")
        if stream:
            resp = await AsyncRequest(
                config=self.config,
                path=path,
                params=cast(Dict[Any, Any], params),
                verb="post",
            ).perform_with_content_streaming()
            return resp
        resp = await AsyncRequest(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp

    async def run(
        self, params: PromptEngineExecuteParams
    ) -> Union[PromptEngineRunResponse, Generator[Any, None, None]]:
        id = params.get("id")
        path = f"/prompt_engine/{id}"
        stream = params.get("stream")

        if stream:
            resp = await AsyncRequest(
                config=self.config,
                path=path,
                params=cast(Dict[Any, Any], params),
                verb="post",
            ).perform_with_content_streaming()
            return resp

        resp = await AsyncRequest(
            config=self.config,
            path=path,
            params=cast(Dict[Any, Any], params),
            verb="post",
        ).perform_with_content()
        return resp