API Reference¶

class AiContext:
    def __init__(
        self,
        client: openai.AsyncClient,
        response_format_capabilities: set[ResponseFormatCapability],
        retry_policy: RetryPolicy | None,
    ) -> None:
        self._async_client = client
        self._response_format_capabilities = response_format_capabilities
        self._retry_policy = _DEFAULT_RETRY_POLICY.merge(retry_policy)

    @property
    def async_client(self) -> openai.AsyncClient:
        return self._async_client

    @property
    def response_format_capabilities(self) -> set[ResponseFormatCapability]:
        return self._response_format_capabilities

    @property
    def retry_policy(self) -> RetryPolicy:
        return self._retry_policy

class RetryPolicy(BaseModel):
    network: RetryParameters | None = None
    validation: RetryParameters | None = None

    def merge(self, child: "RetryPolicy | None") -> "RetryPolicy":
        if not child:
            return self

        return RetryPolicy(
            network=self.network.merge(child.network) if self.network is not None else child.network,
            validation=self.validation.merge(child.validation) if self.validation is not None else child.validation,
        )

class RetryParameters(BaseModel):
    stop_after_attempt: int | None = None
    stop_after_delay: float | None = None

    wait_initial: float | None = None
    wait_max: float | None = None

    def merge(self, child: "RetryParameters | None") -> "RetryParameters":
        if not child:
            return self
        overrides = child.model_dump(exclude_none=True)
        return self.model_copy(update=overrides)

    def to_wait(self) -> tenacity.wait.wait_base:
        wait_initial = self.wait_initial or 0
        wait_max = self.wait_max or sys.maxsize / 2

        return tenacity.wait_exponential_jitter(initial=wait_initial, max=wait_max)

    def to_stop(self) -> tenacity.stop.stop_base:
        if self.stop_after_attempt:
            stop_attempts = tenacity.stop_after_attempt(self.stop_after_attempt)
        else:
            stop_attempts = None

        if self.stop_after_delay:
            stop_delay = tenacity.stop_after_delay(self.stop_after_delay)
        else:
            stop_delay = None

        return tenacity.stop_any(*(x for x in (stop_attempts, stop_delay) if x is not None))

class StructuredAiCall(Generic[TRequest, TResponse]):
    def __init__(
        self,
        ctx: AiContext,
        config: StructuredAiCallConfig,
        formatter: RequestFormatter[TRequest],
        parser: ResponseParser[TResponse],
        retry_policy: RetryPolicy | None,
    ) -> None:
        self._ctx = ctx
        self._config = config
        self._formatter = formatter
        self._parser = parser

        retry_policy = self._ctx.retry_policy.merge(retry_policy)
        retry_policy_validation = retry_policy.validation
        retry_policy_network = retry_policy.network
        if retry_policy_validation is None:
            raise ValueError("No validation retry policy")
        if retry_policy_network is None:
            raise ValueError("No network retry policy")
        self._retry_policy_validation = retry_policy_validation
        self._retry_policy_network = retry_policy_network

    async def __call__(self, inputs: TRequest) -> TResponse:
        response_format = cast(
            ResponseFormat, self._parser.to_response_format(capabilities=self._ctx.response_format_capabilities)
        )

        @tenacity.retry(
            stop=self._retry_policy_validation.to_stop(),
            wait=self._retry_policy_validation.to_wait(),
            retry=tenacity.retry_if_exception_type(ValueError),
            reraise=True,
            before_sleep=tenacity.before_log(_LOGGER, log_level=logging.WARNING),
        )
        @tenacity.retry(
            stop=self._retry_policy_network.to_stop(),
            wait=self._retry_policy_network.to_wait(),
            retry=tenacity.retry_if_exception_type((OSError, OpenAIError)),
            reraise=True,
            before_sleep=tenacity.before_log(_LOGGER, log_level=logging.WARNING),
        )
        async def run_with_retry() -> TResponse:
            result = await self._ctx.async_client.chat.completions.create(
                messages=[
                    {"role": "system", "content": self._config.system_message},
                    {"role": "user", "content": self._formatter.format(inputs)},
                ],
                response_format=response_format,
                model=self._config.parameters.model,
                temperature=self._config.parameters.temperature,
            )
            content = result.choices[0].message.content

            if content is None:
                exc = ValueError("Got empty content from AI provider")
                exc.add_note(str(content))
                raise exc

            return self._parser.validate(content)

        return await run_with_retry()

    @classmethod
    def from_files(
        cls,
        ctx: AiContext,
        path: Path,
        formatter: RequestFormatter[TRequest],
        parser: ResponseParser[TResponse],
        retry_policy: RetryPolicy | None = None,
    ) -> Self:
        system_message = path.with_suffix(".txt").read_text(encoding="utf-8").strip()
        parameters = AiCallParameters.model_validate_json(path.with_suffix(".json").read_text(encoding="utf-8"))
        config = StructuredAiCallConfig(system_message=system_message, parameters=parameters)
        return cls(ctx, config, formatter=formatter, parser=parser, retry_policy=retry_policy)

@dataclass
class StructuredAiCallConfig:
    system_message: str
    parameters: AiCallParameters

class AiCallParameters(BaseModel):
    model: str
    temperature: float | None = None

class RequestFormatter(abc.ABC, Generic[TRequest]):
    @abc.abstractmethod
    def format(self, content: TRequest) -> str: ...

class JsonFormatter(RequestFormatter[TRequest], Generic[TRequest]):
    def __init__(self, indent: int | None = None) -> None:
        self._indent = indent

    def format(self, content: TRequest) -> str:
        return content.model_dump_json(indent=self._indent or 0)

class XmlFormatter(RequestFormatter[TRequest], Generic[TRequest]):
    def format(self, content: TRequest) -> str:
        return cast(str, content.to_xml(encoding="unicode"))

class StrFormatter(RequestFormatter):
    def format(self, content: str) -> str:
        return content

class ResponseParser(abc.ABC, Generic[TResult]):
    @property
    @abc.abstractmethod
    def returns(self) -> type[TResult]: ...

    @abc.abstractmethod
    def validate(self, content: str) -> TResult: ...

    @abc.abstractmethod
    def to_response_format(self, capabilities: set[ResponseFormatCapability]) -> dict | None: ...

class JsonParser(ResponseParser[TResult], Generic[TResult]):
    def __init__(self, model_type: type[TResult], response_format: JsonResponseFormat) -> None:
        self._model_type = model_type
        self._format = response_format

    @property
    def returns(self) -> type[TResult]:
        return self._model_type

    def validate(self, content: str) -> TResult:
        return self._model_type.model_validate_json(content)

    def to_response_format(self, capabilities: set[ResponseFormatCapability]) -> dict | None:
        if self._format == JsonResponseFormat.json_schema and ResponseFormatCapability.json_schema in capabilities:
            return {
                "type": "json_schema",
                "json_schema": {
                    "name": self._model_type.__name__,
                    "strict": True,
                    "schema": _make_schema_strict(self._model_type.model_json_schema()),
                },
            }
        elif (
            self._format in {JsonResponseFormat.json_schema, JsonResponseFormat.json}
            and ResponseFormatCapability.json in capabilities
        ):
            return {"type": "json_object"}

        return None

class XmlParser(ResponseParser[TResult], Generic[TResult]):
    def __init__(self, model_type: type[TResult]) -> None:
        self._model_type = model_type

    @property
    def returns(self) -> type[TResult]:
        return self._model_type

    def validate(self, content: str) -> TResult:
        return self._model_type.from_xml(content)

    def to_response_format(self, capabilities: set[ResponseFormatCapability]) -> dict | None:
        return None

class StrParser(ResponseParser[str]):
    @property
    def returns(self) -> type[str]:
        return str

    def validate(self, content: str) -> str:
        return content

    def to_response_format(self, capabilities: set[ResponseFormatCapability]) -> dict | None:
        return None