API Reference

AnthropicAdapter

Bases: ModelProviderAdapter

Adapter for Anthropic Claude models in AWS Bedrock.

This adapter handles: 1. Setting the required anthropic_version 2. Building tool definitions in Anthropic's format 3. Validating tool-use responses from Claude models

Source code in src/llmbo/adapters/anthropic.py

class AnthropicAdapter(ModelProviderAdapter):
    """Adapter for Anthropic Claude models in AWS Bedrock.

    This adapter handles:
    1. Setting the required anthropic_version
    2. Building tool definitions in Anthropic's format
    3. Validating tool-use responses from Claude models
    """

    logger = logging.getLogger(f"{__name__}.AnthropicAdapter")

    @classmethod
    @lru_cache(maxsize=1)
    def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
        """Build a tool definition in Anthropic's format."""
        cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")
        tool = {
            "name": output_model.__name__,
            "description": output_model.__doc__ or "Please fill in the schema",
            "input_schema": output_model.model_json_schema(),
        }
        cls.logger.debug(f"Created tool definition with name: {tool['name']}")
        return tool

    @classmethod
    def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
        """Prepare model input for Anthropic Claude models."""
        cls.logger.debug("Preparing model input for Anthropic Claude")

        # Ensure anthropic_version is set (required for Anthropic models)
        if model_input.anthropic_version is None:
            cls.logger.debug("Setting default anthropic_version")
            model_input.anthropic_version = "bedrock-2023-05-31"

        # Build tool from output_model and add it to model_input
        if output_model:
            cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
            tool = cls.build_tool(output_model)
            model_input.tools = [tool]
            model_input.tool_choice = ToolChoice(type="tool", name=tool["name"])

        return model_input

    @classmethod
    def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
        """Validate and parse output from Anthropic Claude models."""
        cls.logger.debug(f"Validating result against {output_model.__name__} schema")

        if result.get("stop_reason") != "tool_use":
            cls.logger.debug(f"Invalid stop_reason: {result.get('stop_reason')}")
            return None

        # Ensure content exists
        content = result.get("content", [])
        if not content:
            cls.logger.debug("Result contains no content")
            return None

        # Check that there's exactly one tool call
        tool_use_items = [item for item in content if item.get("type") == "tool_use"]
        if len(tool_use_items) != 1:
            cls.logger.debug(f"Expected exactly 1 tool_use item, found {len(tool_use_items)}")
            return None

        # Process the single tool use response
        tool_use = tool_use_items[0]
        try:
            # Parse tool use input as our output model
            instance = output_model(**tool_use["input"])
            cls.logger.debug(f"Successfully validated result as {output_model.__name__}")
            return instance
        except ValidationError as e:
            cls.logger.debug(f"Validation failed: {e!s}")
            return None

build_tool(output_model) cached classmethod

Build a tool definition in Anthropic's format.

Source code in src/llmbo/adapters/anthropic.py

@classmethod
@lru_cache(maxsize=1)
def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
    """Build a tool definition in Anthropic's format."""
    cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")
    tool = {
        "name": output_model.__name__,
        "description": output_model.__doc__ or "Please fill in the schema",
        "input_schema": output_model.model_json_schema(),
    }
    cls.logger.debug(f"Created tool definition with name: {tool['name']}")
    return tool

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for Anthropic Claude models.

Source code in src/llmbo/adapters/anthropic.py

@classmethod
def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
    """Prepare model input for Anthropic Claude models."""
    cls.logger.debug("Preparing model input for Anthropic Claude")

    # Ensure anthropic_version is set (required for Anthropic models)
    if model_input.anthropic_version is None:
        cls.logger.debug("Setting default anthropic_version")
        model_input.anthropic_version = "bedrock-2023-05-31"

    # Build tool from output_model and add it to model_input
    if output_model:
        cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
        tool = cls.build_tool(output_model)
        model_input.tools = [tool]
        model_input.tool_choice = ToolChoice(type="tool", name=tool["name"])

    return model_input

validate_result(result, output_model) classmethod

Validate and parse output from Anthropic Claude models.

Source code in src/llmbo/adapters/anthropic.py

@classmethod
def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
    """Validate and parse output from Anthropic Claude models."""
    cls.logger.debug(f"Validating result against {output_model.__name__} schema")

    if result.get("stop_reason") != "tool_use":
        cls.logger.debug(f"Invalid stop_reason: {result.get('stop_reason')}")
        return None

    # Ensure content exists
    content = result.get("content", [])
    if not content:
        cls.logger.debug("Result contains no content")
        return None

    # Check that there's exactly one tool call
    tool_use_items = [item for item in content if item.get("type") == "tool_use"]
    if len(tool_use_items) != 1:
        cls.logger.debug(f"Expected exactly 1 tool_use item, found {len(tool_use_items)}")
        return None

    # Process the single tool use response
    tool_use = tool_use_items[0]
    try:
        # Parse tool use input as our output model
        instance = output_model(**tool_use["input"])
        cls.logger.debug(f"Successfully validated result as {output_model.__name__}")
        return instance
    except ValidationError as e:
        cls.logger.debug(f"Validation failed: {e!s}")
        return None

BatchInferer

A class to manage batch inference jobs using AWS Bedrock.

This class handles the creation, monitoring, and retrieval of batch inference jobs for large-scale model invocations using AWS Bedrock service.

Parameters:

Name	Type	Description	Default
model_name	str	The name/ID of the AWS Bedrock model to use	required
bucket_name	str	The S3 bucket name for storing input/output data	required
region	str	The region to run the batch inference job in.	required
job_name	str	A unique name for the batch inference job	required
role_arn	str	The AWS IAM role ARN with necessary permissions	required
time_out_duration_hours	int	Maximum job runtime in hours. Defaults to 24.	24
session	session	A boto3 session to be used for calls to AWS, If one if not provided a new one will be created	None

Attributes:

Name	Type	Description
job_arn	str	The ARN of the created batch inference job
results	List[dict]	The results of the batch inference job. Available after job completion.
manifest	Manifest	Job execution statistics. Available after job completion.
job_status	str	Current status of the batch job. One of VALID_FINISHED_STATUSES.

Source code in src/llmbo/batch_inferer.py

class BatchInferer:
    """A class to manage batch inference jobs using AWS Bedrock.

    This class handles the creation, monitoring, and retrieval of batch inference jobs
    for large-scale model invocations using AWS Bedrock service.

    Args:
        model_name (str): The name/ID of the AWS Bedrock model to use
        bucket_name (str): The S3 bucket name for storing input/output data
        region (str): The region to run the batch inference job in.
        job_name (str): A unique name for the batch inference job
        role_arn (str): The AWS IAM role ARN with necessary permissions
        time_out_duration_hours (int, optional): Maximum job runtime in hours.
            Defaults to 24.
        session (boto3.session, optional): A boto3 session to be used for calls to AWS,
            If one if not provided a new one will be  created

    Attributes:
        job_arn (str): The ARN of the created batch inference job
        results (List[dict]): The results of the batch inference job.
            Available after job completion.
        manifest (Manifest): Job execution statistics. Available after job completion.
        job_status (str): Current status of the batch job.
            One of VALID_FINISHED_STATUSES.
    """

    logger = logging.getLogger(f"{__name__}.BatchInferer")

    def __init__(
        self,
        model_name: str,  # this should be an enum...
        bucket_name: str,
        region: str,
        job_name: str,
        role_arn: str,
        time_out_duration_hours: int = 24,
        session: boto3.Session | None = None,
        output_dir: str = ".",
    ):
        """Initialize a BatchInferer for AWS Bedrock batch processing.

        Creates a configured batch inference manager that handles the end-to-end process
        of submitting and managing batch jobs on AWS Bedrock.

        Args:
            model_name (str): The AWS Bedrock model identifier
                (e.g., 'anthropic.claude-3-haiku-20240307-v1:0')
            bucket_name (str): Name of the S3 bucket for storing job inputs and outputs
            region (str): The region containing the llm to call, must match the bucket
            job_name (str): Unique identifier for this batch job. Used in file naming.
            role_arn (str): AWS IAM role ARN with permissions for Bedrock and S3 access
            time_out_duration_hours (int, optional): Maximum runtime for the batch job.
                Defaults to 24 hours.
            session (boto3.session, optional): A boto3 session to be used for AWS calls,
                If one if not provided a new one will be created
            output_dir (str, optional): Directory for local JSONL files (input, output,
                manifest). Defaults to "." (current working directory). The directory
                is created if it does not exist.

        Raises:
            KeyError: If AWS_PROFILE environment variable is not set
            ValueError: If the provided role_arn doesn't exist or is invalid

        Example:
        ```python
            >>> bi = BatchInferer(
                    model_name="anthropic.claude-3-haiku-20240307-v1:0",
                    bucket_name="my-inference-bucket",
                    job_name="batch-job-2024-01-01",
                    role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole"
                )
        ```

        Note:
            - Requires valid AWS credentials and configuration
            - The S3 bucket must exist and be accessible via the provided role
            - Job name will be used to create unique file names for inputs and outputs
        """
        self.logger.info("Intialising BatchInferer")
        # model parameters
        self.model_name = model_name
        self.adapter = self._get_adapter(model_name)

        self.time_out_duration_hours = time_out_duration_hours

        self.session: boto3.Session = session or boto3.Session()

        # file/bucket parameters
        self._check_bucket(bucket_name, region)
        self.bucket_name = bucket_name
        self.bucket_uri = "s3://" + bucket_name
        self.job_name = job_name or "batch_inference_" + str(uuid4())[:6]
        self.file_name = job_name + ".jsonl"
        self.output_file_name = None
        self.manifest_file_name = None
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)

        self.check_for_profile()
        self._check_arn(role_arn)
        self.role_arn = role_arn
        self.region = region

        self.client: boto3.client = self.session.client("bedrock", region_name=region)

        # internal state - created by the class later.
        self.job_arn = None
        self.job_status = None
        self.results = None
        self.manifest = None
        self.requests = None

        self.logger.info("Initialized BatchInferer")

    @property
    def unique_id_from_arn(self) -> str:
        """Retrieves the id from the job ARN.

        Raises:
            ValueError: if no job ARN has been set

        Returns:
            str: a unique id portion of the job ARN
        """
        if not self.job_arn:
            self.logger.error("Job ARN not set")
            raise ValueError("Job ARN not set")
        return self.job_arn.split("/")[-1]

    def _local_path(self, bare_name: str) -> str:
        """Join output_dir with a bare filename for local file operations."""
        return os.path.join(self.output_dir, bare_name)

    def check_for_profile(self) -> None:
        """Checks if a profile has been set.

        Raises:
            KeyError: If AWS_PROFILE does not exist in the env.
        """
        if not os.getenv("AWS_PROFILE"):
            self.logger.error("AWS_PROFILE environment variable not set")
            raise KeyError("AWS_PROFILE environment variable not set")

    @staticmethod
    def _read_jsonl(file_path):
        data = []
        with open(file_path) as file:
            for line in file:
                data.append(json.loads(line.strip()))
        return data

    def _get_bucket_location(self, bucket_name: str) -> str | None:
        """Get the location of the s3 bucket.

        Args:
            bucket_name (str): the name of a bucket

        Raises:
            ValueError: If the bucket is not accessible

        Returns:
            str: a region, e.g. "eu-west-2"
        """
        try:
            s3_client = self.session.client("s3")
            response = s3_client.get_bucket_location(Bucket=bucket_name)

            if response:
                region = response["LocationConstraint"]
                # aws returns None if the region is us-east-1 otherwise it returns the
                # region
                return region if region else "us-east-1"
        except ClientError as e:
            self.logger.error(f"Bucket {bucket_name} is not accessible: {e}")
            raise ValueError(f"Bucket {bucket_name} is not accessible") from e

    def _check_bucket(self, bucket_name: str, region: str) -> None:
        """Validate if the bucket_name provided exists.

        Args:
            bucket_name (str): the name of a bucket
            region (str): the name of a region

        Raises:
            ValueError: If the bucket is not accessible
            ValueError: If the bucket is not in the same region as the LLM.
        """
        try:
            s3_client = self.session.client("s3")
            s3_client.head_bucket(Bucket=bucket_name)
        except ClientError as e:
            self.logger.error(f"Bucket {bucket_name} is not accessible: {e}")
            raise ValueError(f"Bucket {bucket_name} is not accessible") from e

        if (bucket_region := self._get_bucket_location(bucket_name)) != region:
            self.logger.error(
                f"Bucket {bucket_name} is not located in the same region [{region}] as the llm [{bucket_region}]"
            )
            raise ValueError(
                f"Bucket {bucket_name} is not located in the same region [{region}] as the llm [{bucket_region}]"
            )

    def _check_arn(self, role_arn: str) -> bool:
        """Validate if an IAM role exists and is accessible.

        Attempts to retrieve the IAM role using the provided ARN to verify its
        existence and accessibility.

        Args:
            role_arn (str): The AWS ARN of the IAM role to check.
                Format: 'arn:aws:iam::<account-id>:role/<role-name>'

        Returns:
            bool: True if the role exists and is accessible.

        Raises:
            ValueError: If the role does not exist.
        ClientError: If there are AWS API issues unrelated to role existence.
        """
        if not role_arn.startswith("arn:aws:iam::"):
            self.logger.error("Invalid ARN format")
            raise ValueError("Invalid ARN format")

        # Extract the role name from the ARN
        role_name = role_arn.split("/")[-1]

        iam_client = self.session.client("iam")

        try:
            # Try to get the role
            iam_client.get_role(RoleName=role_name)
            self.logger.info(f"Role '{role_name}' exists.")
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchEntity":
                self.logger.error(f"Role '{role_name}' does not exist.")
                raise ValueError(f"Role '{role_name}' does not exist.") from e
            else:
                raise e

    def _get_adapter(self, model_name):
        # Get the appropriate adapter for this model
        try:
            adapter = ModelAdapterRegistry.get_adapter(model_name)
            self.logger.info(f"Using {adapter.__name__} for model {model_name}")
            return adapter
        except ValueError as e:
            self.logger.error(f"No adapter available for model {model_name}")
            raise ValueError(f"No adapter available for model {model_name}") from e

    def prepare_requests(self, inputs: dict[str, ModelInput]) -> None:
        """Prepare batch inference requests from a dictionary of model inputs.

        Formats model inputs into the required JSONL structure for AWS Bedrock
        batch processing. Each request is formatted as:
            {
                "recordId": str,
                "modelInput": dict
            }

        Args:
            inputs (Dict[str, ModelInput]): Dictionary mapping record IDs to their corresponding
                ModelInput configurations. The record IDs will be used to track results.

        Raises:
            ValueError: If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

        Example:
            >>> inputs = {
            ...     "001": ModelInput(
            ...         messages=[{"role": "user", "content": "Hello"}],
            ...         temperature=0.7
            ...     ),
            ...     "002": ModelInput(
            ...         messages=[{"role": "user", "content": "Hi"}],
            ...         temperature=0.7
            ...     )
            ... }
            >>> bi.prepare_requests(inputs)

        Note:
            - This method must be called before push_requests_to_s3()
            - The prepared requests are stored in self.requests
            - Each ModelInput is converted to a dict using its to_dict() method
        """
        # TODO: Should I copy these inputs so I dont modify them.
        self.logger.info(f"Preparing {len(inputs)} requests")
        self._check_input_length(inputs)
        self.logger.info("Adding model specific parameters to model_input")
        for id, model_input in inputs.items():
            inputs[id] = self.adapter.prepare_model_input(model_input)

        self.requests = self._to_requests(inputs)

    def _to_requests(self, inputs):
        self.logger.info("Converting to dict")
        return [
            {
                "recordId": id,
                "modelInput": model_input.to_dict(),
            }
            for id, model_input in inputs.items()
        ]

    def _check_input_length(self, inputs):
        if inputs is None:
            self.logger.error("Minimum Batch Size is 100, None supplied")
            raise ValueError("Minimum Batch Size is 100, None supplied")

        if len(inputs) < 100:
            self.logger.error(f"Minimum Batch Size is 100, {len(inputs)} given.")
            raise ValueError(f"Minimum Batch Size is 100, {len(inputs)} given.")

    def _write_requests_locally(self) -> None:
        """Write batch inference requests to a local JSONL file.

        Creates or overwrites a local JSONL file containing the prepared inference
        requests. Each line contains a JSON object with recordId and modelInput.

        Raises:
            IOError: If unable to write to the file
            AttributeError: If called before prepare_requests()

        Note:
            - File is named according to self.file_name
            - Internal method used by push_requests_to_s3()
            - Will overwrite existing files with the same name
        """
        self.logger.info(f"Writing {len(self.requests)} requests to {self._local_path(self.file_name)}")
        with open(self._local_path(self.file_name), "w") as file:
            for record in self.requests:
                file.write(json.dumps(record) + "\n")

    def push_requests_to_s3(self) -> dict[str, Any]:
        """Upload batch inference requests to S3.

        Writes the prepared requests to a local JSONL file and uploads it to the
        configured S3 bucket in the 'input/' prefix.

        Returns:
            dict: The S3 upload response from boto3

        Raises:
            IOError: If local file operations fail
            ClientError: If S3 upload fails
            AttributeError: If called before prepare_requests()

        Note:
            - Creates/overwrites files both locally and in S3
            - S3 path: {bucket_name}/input/{job_name}.jsonl
            - Sets Content-Type to 'application/json'
        """
        # do I want to write this file locally? - maybe stream it or write it to
        # temp file instead
        self._write_requests_locally()
        s3_client = self.session.client("s3")
        self.logger.info(f"Pushing {len(self.requests)} requests to {self.bucket_name}")
        response = s3_client.upload_file(
            Filename=self._local_path(self.file_name),
            Bucket=self.bucket_name,
            Key=f"input/{self.file_name}",
            ExtraArgs={"ContentType": "application/json"},
        )
        return response

    def create(self) -> dict[str, Any]:
        """Create a new batch inference job in AWS Bedrock.

        Initializes a new model invocation job using the configured parameters
        and uploaded input data.

        Returns:
            dict: The complete response from the create_model_invocation_job API call

        Raises:
            RuntimeError: If job creation fails
            ClientError: For AWS API errors
            ValueError: If required configurations are missing

        Note:
            - Sets self.job_arn on successful creation
            - Input data must be uploaded to S3 before calling this method
            - Job will timeout after self.time_out_duration_hours
        """
        if self.requests:
            self.logger.info(f"Creating job {self.job_name}")
            response = self.client.create_model_invocation_job(
                jobName=self.job_name,
                roleArn=self.role_arn,
                clientRequestToken="string",
                modelId=self.model_name,
                inputDataConfig={
                    "s3InputDataConfig": {
                        "s3InputFormat": "JSONL",
                        "s3Uri": f"{self.bucket_uri}/input/{self.file_name}",
                    }
                },
                outputDataConfig={
                    "s3OutputDataConfig": {
                        "s3Uri": f"{self.bucket_uri}/output/",
                    }
                },
                timeoutDurationInHours=self.time_out_duration_hours,
                tags=[{"key": "bedrock_batch_inference", "value": self.job_name}],
            )

            if response:
                response_status = response["ResponseMetadata"]["HTTPStatusCode"]
                if response_status == 200:
                    self.logger.info(f"Job {self.job_name} created successfully")
                    self.logger.info(f"Assigned jobArn: {response['jobArn']}")
                    self.job_arn = response["jobArn"]
                    return response
                else:
                    self.logger.error(
                        f"There was an error creating the job {self.job_name}, non 200 response from bedrock"
                    )
                    raise RuntimeError(
                        f"There was an error creating the job {self.job_name}, non 200 response from bedrock"
                    )
            else:
                self.logger.error("There was an error creating the job, no response from bedrock")
                raise RuntimeError("There was an error creating the job, no response from bedrock")
        else:
            self.logger.error("There were no prepared requests")
            raise AttributeError("There were no prepared requests")

    def download_results(self) -> None:
        """Download batch inference results from S3.

        Retrieves both the results and manifest files from S3 once the job
        has completed. Files are downloaded to:
            - {job_name}_out.jsonl: Contains model outputs
            - {job_name}_manifest.jsonl: Contains job statistics

        Raises:
            ClientError: For S3 download failures
            ValueError: If job hasn't completed or job_arn isn't set

        Note:
            - Only downloads if job status is in VALID_FINISHED_STATUSES
            - Files are downloaded to current working directory
            - Existing files will be overwritten
            - Call check_complete() first to ensure job is finished
        """
        if self.check_complete() in VALID_FINISHED_STATUSES:
            file_name_, ext = os.path.splitext(self.file_name)
            self.output_file_name = f"{file_name_}_out{ext}"
            self.manifest_file_name = f"{file_name_}_manifest{ext}"
            self.logger.info(f"Job:{self.job_arn} Complete. Downloading results from {self.bucket_name}")
            s3_client = self.session.client("s3")
            s3_client.download_file(
                Bucket=self.bucket_name,
                Key=f"output/{self.unique_id_from_arn}/{self.file_name}.out",
                Filename=self._local_path(self.output_file_name),
            )
            self.logger.info(f"Downloaded results file to {self._local_path(self.output_file_name)}")

            s3_client.download_file(
                Bucket=self.bucket_name,
                Key=f"output/{self.unique_id_from_arn}/manifest.json.out",
                Filename=self._local_path(self.manifest_file_name),
            )
            self.logger.info(f"Downloaded manifest file to {self._local_path(self.manifest_file_name)}")
        else:
            self.logger.info(f"Job:{self.job_arn} was not marked one of {VALID_FINISHED_STATUSES}, could not download.")

    def load_results(self) -> None:
        """Load batch inference results and manifest from local files.

        Reads and parses the output files downloaded from S3, populating:
            - self.results: List of inference results from the output JSONL file
            - self.manifest: Statistics about the job execution (total records, success/error counts, etc.)

        The method expects two files to exist locally:
            - {job_name}_out.jsonl: Contains the model outputs
            - {job_name}_manifest.jsonl: Contains execution statistics

        Raises:
            FileExistsError: If either the results or manifest files are not found locally

        Note:
            - Must call download_results() before calling this method
            - The manifest provides useful metrics like success rate and token counts
        """
        if os.path.isfile(self._local_path(self.output_file_name)) and os.path.isfile(
            self._local_path(self.manifest_file_name)
        ):
            self.results = self._read_jsonl(self._local_path(self.output_file_name))
            self.manifest = Manifest(**self._read_jsonl(self._local_path(self.manifest_file_name))[0])
        else:
            self.logger.error("Result files do not exist, you may need to call .download_results() first.")
            raise FileExistsError("Result files do not exist, you may need to call .download_results() first.")

    def cancel_batch(self) -> None:
        """Cancel a running batch inference job.

        Attempts to stop the currently running batch inference job identified by self.job_arn.

        Returns:
            None

        Raises:
            RuntimeError: If the job cancellation request fails
            ValueError: If no job_arn is set (i.e., no job has been created)
        """
        if not self.job_arn:
            self.logger.error("No job_arn set - no job to cancel")
            raise ValueError("No job_arn set - no job to cancel")

        response = self.client.stop_model_invocation_job(jobIdentifier=self.job_arn)

        if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
            self.logger.info(f"Job {self.job_name} with id={self.job_arn} was cancelled")
            self.job_status = "Stopped"
        else:
            self.logger.error(
                f"Failed to cancel job {self.job_name}. Status: {response['ResponseMetadata']['HTTPStatusCode']}"
            )
            raise RuntimeError(f"Failed to cancel job {self.job_name}")

    def check_complete(self) -> str | None:
        """Check if the batch inference job has completed.

        Returns:
        str | None: The job status if the job has finished (one of 'Completed', 'Failed',
            'Stopped', or 'Expired'), or None if the job is still in progress.
        """
        if self.job_status not in VALID_FINISHED_STATUSES:
            self.logger.info(f"Checking status of job {self.job_arn}")
            response = self.client.get_model_invocation_job(jobIdentifier=self.job_arn)

            self.job_status = response["status"]
            self.logger.info(f"Job status is {self.job_status}")

            if self.job_status in VALID_FINISHED_STATUSES:
                return self.job_status
            return None
        else:
            self.logger.info(f"Job {self.job_arn} is already {self.job_status}")
            return self.job_status

    def poll_progress(self, poll_interval_seconds: int = 60) -> bool:
        """Polls the progress of a job.

        Args:
            poll_interval_seconds (int, optional): Number of seconds between checks. Defaults to 60.

        Returns:
            bool: True if job is complete.
        """
        self.logger.info(f"Polling for progress every {poll_interval_seconds} seconds")
        while not self.check_complete():
            time.sleep(poll_interval_seconds)
        return True

    def auto(self, inputs: dict[str, ModelInput], poll_time_secs: int = 60) -> dict:
        """Execute the complete batch inference workflow automatically.

        This method combines the preparation, execution, monitoring, and result retrieval
        steps into a single operation.

        Args:
            inputs (Dict[str, ModelInput]): Dictionary of record IDs mapped to their ModelInput configurations
            poll_time_secs (int, optional): How often to poll for model progress. Defaults to 60.

        Returns:
            List[Dict]: The results from the batch inference job
        """
        self.prepare_requests(inputs)
        self.push_requests_to_s3()
        self.create()
        self.poll_progress(poll_time_secs)
        self.download_results()
        self.load_results()
        return self.results

    @classmethod
    def recover_details_from_job_arn(
        cls,
        job_arn: str,
        region: str,
        session: boto3.Session | None = None,
        output_dir: str = ".",
    ) -> "BatchInferer":
        """Recover a BatchInferer instance from an existing job ARN.

        Used to reconstruct a BatchInferer object when the original Python process
        has terminated but the AWS job is still running or complete.

        Args:
            job_arn: (str) The AWS ARN of the existing batch inference job
            region: (str) the region where the job was scheduled
            session (boto3.session, optional): A boto3 session to be used for calls to AWS,
                    If one if not provided a new one will be  created
            output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

        Returns:
            BatchInferer: A configured instance with the job's details

        Raises:
            ValueError: If the job cannot be found or response is invalid

        Example:
            >>> job_arn = "arn:aws:bedrock:region:account:job/xyz123"
            >>> bi = BatchInferer.recover_details_from_job_arn(job_arn)
            >>> bi.check_complete()
            'Completed'
        """
        cls.logger.info(f"Attempting to Recover BatchInferer from {job_arn}")
        session = session or boto3.Session()
        response = cls.check_for_existing_job(job_arn, region, session)

        try:
            # Extract required parameters from response
            job_name = response["jobName"]
            model_id = response["modelId"]
            bucket_name = response["inputDataConfig"]["s3InputDataConfig"]["s3Uri"].split("/")[2]
            role_arn = response["roleArn"]

            # Validate required files exist
            input_file = os.path.join(output_dir, f"{job_name}.jsonl")
            if not os.path.exists(input_file):
                cls.logger.error(f"Required input file not found: {input_file}")
                raise FileNotFoundError(f"Required input file not found: {input_file}")

            requests = cls._read_jsonl(input_file)

            bi = cls(
                model_name=model_id,
                job_name=job_name,
                region=region,
                bucket_name=bucket_name,
                role_arn=role_arn,
                session=session,
                output_dir=output_dir,
            )
            bi.job_arn = job_arn
            bi.requests = requests
            bi.job_status = response["status"]

            return bi

        except (KeyError, IndexError) as e:
            cls.logger.error(f"Invalid job response format: {e!s}")
            raise ValueError(f"Invalid job response format: {e!s}") from e
        except Exception as e:
            cls.logger.error(f"Failed to recover job details: {e!s}")
            raise RuntimeError(f"Failed to recover job details: {e!s}") from e

    @classmethod
    def check_for_existing_job(cls, job_arn, region, session: boto3.Session | None = None) -> dict[str, Any]:
        """Check if a job exists and return its details.

        Args:
            job_arn (str): The AWS ARN of the job to check
            region (str): The AWS region where the job was created
            session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                                           If not provided, a new session will be created.

        Returns:
            Dict[str, Any]: The job details from AWS Bedrock

        Raises:
            ValueError: If the job ARN is invalid or the job is not found
            RuntimeError: For other AWS API errors
        """
        if not job_arn.startswith("arn:aws:bedrock:"):
            cls.logger.error(f"Invalid Bedrock ARN format: {job_arn}")
            raise ValueError(f"Invalid Bedrock ARN format: {job_arn}")
        session = session or boto3.Session()
        client = session.client("bedrock", region_name=region)

        try:
            response = client.get_model_invocation_job(jobIdentifier=job_arn)
        except ClientError as e:
            if e.response["Error"]["Code"] == "ResourceNotFoundException":
                cls.logger.error(f"Job not found: {job_arn}")
                raise ValueError(f"Job not found: {job_arn}") from e
            cls.logger.error(f"AWS API error: {e!s}")
            raise RuntimeError(f"AWS API error: {e!s}") from e

        if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
            cls.logger.error(f"Unexpected response status: {response['ResponseMetadata']['HTTPStatusCode']}")
            raise RuntimeError(f"Unexpected response status: {response['ResponseMetadata']['HTTPStatusCode']}")

        return response

unique_id_from_arn property

Retrieves the id from the job ARN.

Raises:

Type	Description
ValueError	if no job ARN has been set

Returns:

Name	Type	Description
str	str	a unique id portion of the job ARN

__init__(model_name, bucket_name, region, job_name, role_arn, time_out_duration_hours=24, session=None, output_dir='.')

Initialize a BatchInferer for AWS Bedrock batch processing.

Creates a configured batch inference manager that handles the end-to-end process of submitting and managing batch jobs on AWS Bedrock.

Parameters:

Name	Type	Description	Default
model_name	str	The AWS Bedrock model identifier (e.g., 'anthropic.claude-3-haiku-20240307-v1:0')	required
bucket_name	str	Name of the S3 bucket for storing job inputs and outputs	required
region	str	The region containing the llm to call, must match the bucket	required
job_name	str	Unique identifier for this batch job. Used in file naming.	required
role_arn	str	AWS IAM role ARN with permissions for Bedrock and S3 access	required
time_out_duration_hours	int	Maximum runtime for the batch job. Defaults to 24 hours.	24
session	session	A boto3 session to be used for AWS calls, If one if not provided a new one will be created	None
output_dir	str	Directory for local JSONL files (input, output, manifest). Defaults to "." (current working directory). The directory is created if it does not exist.	'.'

Raises:

Type	Description
KeyError	If AWS_PROFILE environment variable is not set
ValueError	If the provided role_arn doesn't exist or is invalid

Example:

    >>> bi = BatchInferer(
            model_name="anthropic.claude-3-haiku-20240307-v1:0",
            bucket_name="my-inference-bucket",
            job_name="batch-job-2024-01-01",
            role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole"
        )

Note

• Requires valid AWS credentials and configuration
• The S3 bucket must exist and be accessible via the provided role
• Job name will be used to create unique file names for inputs and outputs

Source code in src/llmbo/batch_inferer.py

def __init__(
    self,
    model_name: str,  # this should be an enum...
    bucket_name: str,
    region: str,
    job_name: str,
    role_arn: str,
    time_out_duration_hours: int = 24,
    session: boto3.Session | None = None,
    output_dir: str = ".",
):
    """Initialize a BatchInferer for AWS Bedrock batch processing.

    Creates a configured batch inference manager that handles the end-to-end process
    of submitting and managing batch jobs on AWS Bedrock.

    Args:
        model_name (str): The AWS Bedrock model identifier
            (e.g., 'anthropic.claude-3-haiku-20240307-v1:0')
        bucket_name (str): Name of the S3 bucket for storing job inputs and outputs
        region (str): The region containing the llm to call, must match the bucket
        job_name (str): Unique identifier for this batch job. Used in file naming.
        role_arn (str): AWS IAM role ARN with permissions for Bedrock and S3 access
        time_out_duration_hours (int, optional): Maximum runtime for the batch job.
            Defaults to 24 hours.
        session (boto3.session, optional): A boto3 session to be used for AWS calls,
            If one if not provided a new one will be created
        output_dir (str, optional): Directory for local JSONL files (input, output,
            manifest). Defaults to "." (current working directory). The directory
            is created if it does not exist.

    Raises:
        KeyError: If AWS_PROFILE environment variable is not set
        ValueError: If the provided role_arn doesn't exist or is invalid

    Example:
    ```python
        >>> bi = BatchInferer(
                model_name="anthropic.claude-3-haiku-20240307-v1:0",
                bucket_name="my-inference-bucket",
                job_name="batch-job-2024-01-01",
                role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole"
            )
    ```

    Note:
        - Requires valid AWS credentials and configuration
        - The S3 bucket must exist and be accessible via the provided role
        - Job name will be used to create unique file names for inputs and outputs
    """
    self.logger.info("Intialising BatchInferer")
    # model parameters
    self.model_name = model_name
    self.adapter = self._get_adapter(model_name)

    self.time_out_duration_hours = time_out_duration_hours

    self.session: boto3.Session = session or boto3.Session()

    # file/bucket parameters
    self._check_bucket(bucket_name, region)
    self.bucket_name = bucket_name
    self.bucket_uri = "s3://" + bucket_name
    self.job_name = job_name or "batch_inference_" + str(uuid4())[:6]
    self.file_name = job_name + ".jsonl"
    self.output_file_name = None
    self.manifest_file_name = None
    self.output_dir = output_dir
    os.makedirs(self.output_dir, exist_ok=True)

    self.check_for_profile()
    self._check_arn(role_arn)
    self.role_arn = role_arn
    self.region = region

    self.client: boto3.client = self.session.client("bedrock", region_name=region)

    # internal state - created by the class later.
    self.job_arn = None
    self.job_status = None
    self.results = None
    self.manifest = None
    self.requests = None

    self.logger.info("Initialized BatchInferer")

_check_arn(role_arn)

Validate if an IAM role exists and is accessible.

Attempts to retrieve the IAM role using the provided ARN to verify its existence and accessibility.

Parameters:

Name	Type	Description	Default
role_arn	str	The AWS ARN of the IAM role to check. Format: 'arn:aws:iam:::role/'	required

Returns:

Name	Type	Description
bool	bool	True if the role exists and is accessible.

Raises:

Type	Description
ValueError	If the role does not exist.

ClientError: If there are AWS API issues unrelated to role existence.

Source code in src/llmbo/batch_inferer.py

def _check_arn(self, role_arn: str) -> bool:
    """Validate if an IAM role exists and is accessible.

    Attempts to retrieve the IAM role using the provided ARN to verify its
    existence and accessibility.

    Args:
        role_arn (str): The AWS ARN of the IAM role to check.
            Format: 'arn:aws:iam::<account-id>:role/<role-name>'

    Returns:
        bool: True if the role exists and is accessible.

    Raises:
        ValueError: If the role does not exist.
    ClientError: If there are AWS API issues unrelated to role existence.
    """
    if not role_arn.startswith("arn:aws:iam::"):
        self.logger.error("Invalid ARN format")
        raise ValueError("Invalid ARN format")

    # Extract the role name from the ARN
    role_name = role_arn.split("/")[-1]

    iam_client = self.session.client("iam")

    try:
        # Try to get the role
        iam_client.get_role(RoleName=role_name)
        self.logger.info(f"Role '{role_name}' exists.")
        return True
    except ClientError as e:
        if e.response["Error"]["Code"] == "NoSuchEntity":
            self.logger.error(f"Role '{role_name}' does not exist.")
            raise ValueError(f"Role '{role_name}' does not exist.") from e
        else:
            raise e

_check_bucket(bucket_name, region)

Validate if the bucket_name provided exists.

Parameters:

Name	Type	Description	Default
bucket_name	str	the name of a bucket	required
region	str	the name of a region	required

Raises:

Type	Description
ValueError	If the bucket is not accessible
ValueError	If the bucket is not in the same region as the LLM.

Source code in src/llmbo/batch_inferer.py

def _check_bucket(self, bucket_name: str, region: str) -> None:
    """Validate if the bucket_name provided exists.

    Args:
        bucket_name (str): the name of a bucket
        region (str): the name of a region

    Raises:
        ValueError: If the bucket is not accessible
        ValueError: If the bucket is not in the same region as the LLM.
    """
    try:
        s3_client = self.session.client("s3")
        s3_client.head_bucket(Bucket=bucket_name)
    except ClientError as e:
        self.logger.error(f"Bucket {bucket_name} is not accessible: {e}")
        raise ValueError(f"Bucket {bucket_name} is not accessible") from e

    if (bucket_region := self._get_bucket_location(bucket_name)) != region:
        self.logger.error(
            f"Bucket {bucket_name} is not located in the same region [{region}] as the llm [{bucket_region}]"
        )
        raise ValueError(
            f"Bucket {bucket_name} is not located in the same region [{region}] as the llm [{bucket_region}]"
        )

_get_bucket_location(bucket_name)

Get the location of the s3 bucket.

Parameters:

Name	Type	Description	Default
bucket_name	str	the name of a bucket	required

Raises:

Type	Description
ValueError	If the bucket is not accessible

Returns:

Name	Type	Description
str	str | None	a region, e.g. "eu-west-2"

Source code in src/llmbo/batch_inferer.py

def _get_bucket_location(self, bucket_name: str) -> str | None:
    """Get the location of the s3 bucket.

    Args:
        bucket_name (str): the name of a bucket

    Raises:
        ValueError: If the bucket is not accessible

    Returns:
        str: a region, e.g. "eu-west-2"
    """
    try:
        s3_client = self.session.client("s3")
        response = s3_client.get_bucket_location(Bucket=bucket_name)

        if response:
            region = response["LocationConstraint"]
            # aws returns None if the region is us-east-1 otherwise it returns the
            # region
            return region if region else "us-east-1"
    except ClientError as e:
        self.logger.error(f"Bucket {bucket_name} is not accessible: {e}")
        raise ValueError(f"Bucket {bucket_name} is not accessible") from e

_local_path(bare_name)

Join output_dir with a bare filename for local file operations.

Source code in src/llmbo/batch_inferer.py

def _local_path(self, bare_name: str) -> str:
    """Join output_dir with a bare filename for local file operations."""
    return os.path.join(self.output_dir, bare_name)

_write_requests_locally()

Write batch inference requests to a local JSONL file.

Creates or overwrites a local JSONL file containing the prepared inference requests. Each line contains a JSON object with recordId and modelInput.

Raises:

Type	Description
IOError	If unable to write to the file
AttributeError	If called before prepare_requests()

Note

• File is named according to self.file_name
• Internal method used by push_requests_to_s3()
• Will overwrite existing files with the same name

Source code in src/llmbo/batch_inferer.py

def _write_requests_locally(self) -> None:
    """Write batch inference requests to a local JSONL file.

    Creates or overwrites a local JSONL file containing the prepared inference
    requests. Each line contains a JSON object with recordId and modelInput.

    Raises:
        IOError: If unable to write to the file
        AttributeError: If called before prepare_requests()

    Note:
        - File is named according to self.file_name
        - Internal method used by push_requests_to_s3()
        - Will overwrite existing files with the same name
    """
    self.logger.info(f"Writing {len(self.requests)} requests to {self._local_path(self.file_name)}")
    with open(self._local_path(self.file_name), "w") as file:
        for record in self.requests:
            file.write(json.dumps(record) + "\n")

auto(inputs, poll_time_secs=60)

Execute the complete batch inference workflow automatically.

This method combines the preparation, execution, monitoring, and result retrieval steps into a single operation.

Parameters:

Name	Type	Description	Default
inputs	Dict[str, ModelInput]	Dictionary of record IDs mapped to their ModelInput configurations	required
poll_time_secs	int	How often to poll for model progress. Defaults to 60.	60

Returns:

Type	Description
dict	List[Dict]: The results from the batch inference job

Source code in src/llmbo/batch_inferer.py

def auto(self, inputs: dict[str, ModelInput], poll_time_secs: int = 60) -> dict:
    """Execute the complete batch inference workflow automatically.

    This method combines the preparation, execution, monitoring, and result retrieval
    steps into a single operation.

    Args:
        inputs (Dict[str, ModelInput]): Dictionary of record IDs mapped to their ModelInput configurations
        poll_time_secs (int, optional): How often to poll for model progress. Defaults to 60.

    Returns:
        List[Dict]: The results from the batch inference job
    """
    self.prepare_requests(inputs)
    self.push_requests_to_s3()
    self.create()
    self.poll_progress(poll_time_secs)
    self.download_results()
    self.load_results()
    return self.results

cancel_batch()

Cancel a running batch inference job.

Attempts to stop the currently running batch inference job identified by self.job_arn.

Returns:

Type	Description
None	None

Raises:

Type	Description
RuntimeError	If the job cancellation request fails
ValueError	If no job_arn is set (i.e., no job has been created)

Source code in src/llmbo/batch_inferer.py

def cancel_batch(self) -> None:
    """Cancel a running batch inference job.

    Attempts to stop the currently running batch inference job identified by self.job_arn.

    Returns:
        None

    Raises:
        RuntimeError: If the job cancellation request fails
        ValueError: If no job_arn is set (i.e., no job has been created)
    """
    if not self.job_arn:
        self.logger.error("No job_arn set - no job to cancel")
        raise ValueError("No job_arn set - no job to cancel")

    response = self.client.stop_model_invocation_job(jobIdentifier=self.job_arn)

    if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
        self.logger.info(f"Job {self.job_name} with id={self.job_arn} was cancelled")
        self.job_status = "Stopped"
    else:
        self.logger.error(
            f"Failed to cancel job {self.job_name}. Status: {response['ResponseMetadata']['HTTPStatusCode']}"
        )
        raise RuntimeError(f"Failed to cancel job {self.job_name}")

check_complete()

Check if the batch inference job has completed.

str | None: The job status if the job has finished (one of 'Completed', 'Failed', 'Stopped', or 'Expired'), or None if the job is still in progress.

Source code in src/llmbo/batch_inferer.py

def check_complete(self) -> str | None:
    """Check if the batch inference job has completed.

    Returns:
    str | None: The job status if the job has finished (one of 'Completed', 'Failed',
        'Stopped', or 'Expired'), or None if the job is still in progress.
    """
    if self.job_status not in VALID_FINISHED_STATUSES:
        self.logger.info(f"Checking status of job {self.job_arn}")
        response = self.client.get_model_invocation_job(jobIdentifier=self.job_arn)

        self.job_status = response["status"]
        self.logger.info(f"Job status is {self.job_status}")

        if self.job_status in VALID_FINISHED_STATUSES:
            return self.job_status
        return None
    else:
        self.logger.info(f"Job {self.job_arn} is already {self.job_status}")
        return self.job_status

check_for_existing_job(job_arn, region, session=None) classmethod

Check if a job exists and return its details.

Parameters:

Name	Type	Description	Default
job_arn	str	The AWS ARN of the job to check	required
region	str	The AWS region where the job was created	required
session	Session	A boto3 session to be used for AWS API calls. If not provided, a new session will be created.	None

Returns:

Type	Description
dict[str, Any]	Dict[str, Any]: The job details from AWS Bedrock

Raises:

Type	Description
ValueError	If the job ARN is invalid or the job is not found
RuntimeError	For other AWS API errors

Source code in src/llmbo/batch_inferer.py

@classmethod
def check_for_existing_job(cls, job_arn, region, session: boto3.Session | None = None) -> dict[str, Any]:
    """Check if a job exists and return its details.

    Args:
        job_arn (str): The AWS ARN of the job to check
        region (str): The AWS region where the job was created
        session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                                       If not provided, a new session will be created.

    Returns:
        Dict[str, Any]: The job details from AWS Bedrock

    Raises:
        ValueError: If the job ARN is invalid or the job is not found
        RuntimeError: For other AWS API errors
    """
    if not job_arn.startswith("arn:aws:bedrock:"):
        cls.logger.error(f"Invalid Bedrock ARN format: {job_arn}")
        raise ValueError(f"Invalid Bedrock ARN format: {job_arn}")
    session = session or boto3.Session()
    client = session.client("bedrock", region_name=region)

    try:
        response = client.get_model_invocation_job(jobIdentifier=job_arn)
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceNotFoundException":
            cls.logger.error(f"Job not found: {job_arn}")
            raise ValueError(f"Job not found: {job_arn}") from e
        cls.logger.error(f"AWS API error: {e!s}")
        raise RuntimeError(f"AWS API error: {e!s}") from e

    if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
        cls.logger.error(f"Unexpected response status: {response['ResponseMetadata']['HTTPStatusCode']}")
        raise RuntimeError(f"Unexpected response status: {response['ResponseMetadata']['HTTPStatusCode']}")

    return response

check_for_profile()

Checks if a profile has been set.

Raises:

Type	Description
KeyError	If AWS_PROFILE does not exist in the env.

Source code in src/llmbo/batch_inferer.py

def check_for_profile(self) -> None:
    """Checks if a profile has been set.

    Raises:
        KeyError: If AWS_PROFILE does not exist in the env.
    """
    if not os.getenv("AWS_PROFILE"):
        self.logger.error("AWS_PROFILE environment variable not set")
        raise KeyError("AWS_PROFILE environment variable not set")

create()

Create a new batch inference job in AWS Bedrock.

Initializes a new model invocation job using the configured parameters and uploaded input data.

Returns:

Name	Type	Description
dict	dict[str, Any]	The complete response from the create_model_invocation_job API call

Raises:

Type	Description
RuntimeError	If job creation fails
ClientError	For AWS API errors
ValueError	If required configurations are missing

Note

• Sets self.job_arn on successful creation
• Input data must be uploaded to S3 before calling this method
• Job will timeout after self.time_out_duration_hours

Source code in src/llmbo/batch_inferer.py

def create(self) -> dict[str, Any]:
    """Create a new batch inference job in AWS Bedrock.

    Initializes a new model invocation job using the configured parameters
    and uploaded input data.

    Returns:
        dict: The complete response from the create_model_invocation_job API call

    Raises:
        RuntimeError: If job creation fails
        ClientError: For AWS API errors
        ValueError: If required configurations are missing

    Note:
        - Sets self.job_arn on successful creation
        - Input data must be uploaded to S3 before calling this method
        - Job will timeout after self.time_out_duration_hours
    """
    if self.requests:
        self.logger.info(f"Creating job {self.job_name}")
        response = self.client.create_model_invocation_job(
            jobName=self.job_name,
            roleArn=self.role_arn,
            clientRequestToken="string",
            modelId=self.model_name,
            inputDataConfig={
                "s3InputDataConfig": {
                    "s3InputFormat": "JSONL",
                    "s3Uri": f"{self.bucket_uri}/input/{self.file_name}",
                }
            },
            outputDataConfig={
                "s3OutputDataConfig": {
                    "s3Uri": f"{self.bucket_uri}/output/",
                }
            },
            timeoutDurationInHours=self.time_out_duration_hours,
            tags=[{"key": "bedrock_batch_inference", "value": self.job_name}],
        )

        if response:
            response_status = response["ResponseMetadata"]["HTTPStatusCode"]
            if response_status == 200:
                self.logger.info(f"Job {self.job_name} created successfully")
                self.logger.info(f"Assigned jobArn: {response['jobArn']}")
                self.job_arn = response["jobArn"]
                return response
            else:
                self.logger.error(
                    f"There was an error creating the job {self.job_name}, non 200 response from bedrock"
                )
                raise RuntimeError(
                    f"There was an error creating the job {self.job_name}, non 200 response from bedrock"
                )
        else:
            self.logger.error("There was an error creating the job, no response from bedrock")
            raise RuntimeError("There was an error creating the job, no response from bedrock")
    else:
        self.logger.error("There were no prepared requests")
        raise AttributeError("There were no prepared requests")

download_results()

Download batch inference results from S3.

Retrieves both the results and manifest files from S3 once the job has completed. Files are downloaded to: - {job_name}_out.jsonl: Contains model outputs - {job_name}_manifest.jsonl: Contains job statistics

Raises:

Type	Description
ClientError	For S3 download failures
ValueError	If job hasn't completed or job_arn isn't set

Note

• Only downloads if job status is in VALID_FINISHED_STATUSES
• Files are downloaded to current working directory
• Existing files will be overwritten
• Call check_complete() first to ensure job is finished

Source code in src/llmbo/batch_inferer.py

def download_results(self) -> None:
    """Download batch inference results from S3.

    Retrieves both the results and manifest files from S3 once the job
    has completed. Files are downloaded to:
        - {job_name}_out.jsonl: Contains model outputs
        - {job_name}_manifest.jsonl: Contains job statistics

    Raises:
        ClientError: For S3 download failures
        ValueError: If job hasn't completed or job_arn isn't set

    Note:
        - Only downloads if job status is in VALID_FINISHED_STATUSES
        - Files are downloaded to current working directory
        - Existing files will be overwritten
        - Call check_complete() first to ensure job is finished
    """
    if self.check_complete() in VALID_FINISHED_STATUSES:
        file_name_, ext = os.path.splitext(self.file_name)
        self.output_file_name = f"{file_name_}_out{ext}"
        self.manifest_file_name = f"{file_name_}_manifest{ext}"
        self.logger.info(f"Job:{self.job_arn} Complete. Downloading results from {self.bucket_name}")
        s3_client = self.session.client("s3")
        s3_client.download_file(
            Bucket=self.bucket_name,
            Key=f"output/{self.unique_id_from_arn}/{self.file_name}.out",
            Filename=self._local_path(self.output_file_name),
        )
        self.logger.info(f"Downloaded results file to {self._local_path(self.output_file_name)}")

        s3_client.download_file(
            Bucket=self.bucket_name,
            Key=f"output/{self.unique_id_from_arn}/manifest.json.out",
            Filename=self._local_path(self.manifest_file_name),
        )
        self.logger.info(f"Downloaded manifest file to {self._local_path(self.manifest_file_name)}")
    else:
        self.logger.info(f"Job:{self.job_arn} was not marked one of {VALID_FINISHED_STATUSES}, could not download.")

load_results()

Load batch inference results and manifest from local files.

Reads and parses the output files downloaded from S3, populating: - self.results: List of inference results from the output JSONL file - self.manifest: Statistics about the job execution (total records, success/error counts, etc.)

The method expects two files to exist locally

• {job_name}_out.jsonl: Contains the model outputs
• {job_name}_manifest.jsonl: Contains execution statistics

Raises:

Type	Description
FileExistsError	If either the results or manifest files are not found locally

Note

• Must call download_results() before calling this method
• The manifest provides useful metrics like success rate and token counts

Source code in src/llmbo/batch_inferer.py

def load_results(self) -> None:
    """Load batch inference results and manifest from local files.

    Reads and parses the output files downloaded from S3, populating:
        - self.results: List of inference results from the output JSONL file
        - self.manifest: Statistics about the job execution (total records, success/error counts, etc.)

    The method expects two files to exist locally:
        - {job_name}_out.jsonl: Contains the model outputs
        - {job_name}_manifest.jsonl: Contains execution statistics

    Raises:
        FileExistsError: If either the results or manifest files are not found locally

    Note:
        - Must call download_results() before calling this method
        - The manifest provides useful metrics like success rate and token counts
    """
    if os.path.isfile(self._local_path(self.output_file_name)) and os.path.isfile(
        self._local_path(self.manifest_file_name)
    ):
        self.results = self._read_jsonl(self._local_path(self.output_file_name))
        self.manifest = Manifest(**self._read_jsonl(self._local_path(self.manifest_file_name))[0])
    else:
        self.logger.error("Result files do not exist, you may need to call .download_results() first.")
        raise FileExistsError("Result files do not exist, you may need to call .download_results() first.")

poll_progress(poll_interval_seconds=60)

Polls the progress of a job.

Parameters:

Name	Type	Description	Default
poll_interval_seconds	int	Number of seconds between checks. Defaults to 60.	60

Returns:

Name	Type	Description
bool	bool	True if job is complete.

Source code in src/llmbo/batch_inferer.py

def poll_progress(self, poll_interval_seconds: int = 60) -> bool:
    """Polls the progress of a job.

    Args:
        poll_interval_seconds (int, optional): Number of seconds between checks. Defaults to 60.

    Returns:
        bool: True if job is complete.
    """
    self.logger.info(f"Polling for progress every {poll_interval_seconds} seconds")
    while not self.check_complete():
        time.sleep(poll_interval_seconds)
    return True

prepare_requests(inputs)

Prepare batch inference requests from a dictionary of model inputs.

Formats model inputs into the required JSONL structure for AWS Bedrock batch processing. Each request is formatted as: { "recordId": str, "modelInput": dict }

Parameters:

Name	Type	Description	Default
inputs	Dict[str, ModelInput]	Dictionary mapping record IDs to their corresponding ModelInput configurations. The record IDs will be used to track results.	required

Raises:

Type	Description
ValueError	If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

Example

inputs = { ... "001": ModelInput( ... messages=[{"role": "user", "content": "Hello"}], ... temperature=0.7 ... ), ... "002": ModelInput( ... messages=[{"role": "user", "content": "Hi"}], ... temperature=0.7 ... ) ... } bi.prepare_requests(inputs)

Note

• This method must be called before push_requests_to_s3()
• The prepared requests are stored in self.requests
• Each ModelInput is converted to a dict using its to_dict() method

Source code in src/llmbo/batch_inferer.py

def prepare_requests(self, inputs: dict[str, ModelInput]) -> None:
    """Prepare batch inference requests from a dictionary of model inputs.

    Formats model inputs into the required JSONL structure for AWS Bedrock
    batch processing. Each request is formatted as:
        {
            "recordId": str,
            "modelInput": dict
        }

    Args:
        inputs (Dict[str, ModelInput]): Dictionary mapping record IDs to their corresponding
            ModelInput configurations. The record IDs will be used to track results.

    Raises:
        ValueError: If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

    Example:
        >>> inputs = {
        ...     "001": ModelInput(
        ...         messages=[{"role": "user", "content": "Hello"}],
        ...         temperature=0.7
        ...     ),
        ...     "002": ModelInput(
        ...         messages=[{"role": "user", "content": "Hi"}],
        ...         temperature=0.7
        ...     )
        ... }
        >>> bi.prepare_requests(inputs)

    Note:
        - This method must be called before push_requests_to_s3()
        - The prepared requests are stored in self.requests
        - Each ModelInput is converted to a dict using its to_dict() method
    """
    # TODO: Should I copy these inputs so I dont modify them.
    self.logger.info(f"Preparing {len(inputs)} requests")
    self._check_input_length(inputs)
    self.logger.info("Adding model specific parameters to model_input")
    for id, model_input in inputs.items():
        inputs[id] = self.adapter.prepare_model_input(model_input)

    self.requests = self._to_requests(inputs)

push_requests_to_s3()

Upload batch inference requests to S3.

Writes the prepared requests to a local JSONL file and uploads it to the configured S3 bucket in the 'input/' prefix.

Returns:

Name	Type	Description
dict	dict[str, Any]	The S3 upload response from boto3

Raises:

Type	Description
IOError	If local file operations fail
ClientError	If S3 upload fails
AttributeError	If called before prepare_requests()

Note

• Creates/overwrites files both locally and in S3
• S3 path: {bucket_name}/input/{job_name}.jsonl
• Sets Content-Type to 'application/json'

Source code in src/llmbo/batch_inferer.py

def push_requests_to_s3(self) -> dict[str, Any]:
    """Upload batch inference requests to S3.

    Writes the prepared requests to a local JSONL file and uploads it to the
    configured S3 bucket in the 'input/' prefix.

    Returns:
        dict: The S3 upload response from boto3

    Raises:
        IOError: If local file operations fail
        ClientError: If S3 upload fails
        AttributeError: If called before prepare_requests()

    Note:
        - Creates/overwrites files both locally and in S3
        - S3 path: {bucket_name}/input/{job_name}.jsonl
        - Sets Content-Type to 'application/json'
    """
    # do I want to write this file locally? - maybe stream it or write it to
    # temp file instead
    self._write_requests_locally()
    s3_client = self.session.client("s3")
    self.logger.info(f"Pushing {len(self.requests)} requests to {self.bucket_name}")
    response = s3_client.upload_file(
        Filename=self._local_path(self.file_name),
        Bucket=self.bucket_name,
        Key=f"input/{self.file_name}",
        ExtraArgs={"ContentType": "application/json"},
    )
    return response

recover_details_from_job_arn(job_arn, region, session=None, output_dir='.') classmethod

Recover a BatchInferer instance from an existing job ARN.

Used to reconstruct a BatchInferer object when the original Python process has terminated but the AWS job is still running or complete.

Parameters:

Name	Type	Description	Default
job_arn	str	(str) The AWS ARN of the existing batch inference job	required
region	str	(str) the region where the job was scheduled	required
session	session	A boto3 session to be used for calls to AWS, If one if not provided a new one will be created	None
output_dir	str	Directory for local JSONL files. Defaults to ".".	'.'

Returns:

Name	Type	Description
BatchInferer	BatchInferer	A configured instance with the job's details

Raises:

Type	Description
ValueError	If the job cannot be found or response is invalid

Example

job_arn = "arn:aws:bedrock:region:account:job/xyz123" bi = BatchInferer.recover_details_from_job_arn(job_arn) bi.check_complete() 'Completed'

Source code in src/llmbo/batch_inferer.py

@classmethod
def recover_details_from_job_arn(
    cls,
    job_arn: str,
    region: str,
    session: boto3.Session | None = None,
    output_dir: str = ".",
) -> "BatchInferer":
    """Recover a BatchInferer instance from an existing job ARN.

    Used to reconstruct a BatchInferer object when the original Python process
    has terminated but the AWS job is still running or complete.

    Args:
        job_arn: (str) The AWS ARN of the existing batch inference job
        region: (str) the region where the job was scheduled
        session (boto3.session, optional): A boto3 session to be used for calls to AWS,
                If one if not provided a new one will be  created
        output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

    Returns:
        BatchInferer: A configured instance with the job's details

    Raises:
        ValueError: If the job cannot be found or response is invalid

    Example:
        >>> job_arn = "arn:aws:bedrock:region:account:job/xyz123"
        >>> bi = BatchInferer.recover_details_from_job_arn(job_arn)
        >>> bi.check_complete()
        'Completed'
    """
    cls.logger.info(f"Attempting to Recover BatchInferer from {job_arn}")
    session = session or boto3.Session()
    response = cls.check_for_existing_job(job_arn, region, session)

    try:
        # Extract required parameters from response
        job_name = response["jobName"]
        model_id = response["modelId"]
        bucket_name = response["inputDataConfig"]["s3InputDataConfig"]["s3Uri"].split("/")[2]
        role_arn = response["roleArn"]

        # Validate required files exist
        input_file = os.path.join(output_dir, f"{job_name}.jsonl")
        if not os.path.exists(input_file):
            cls.logger.error(f"Required input file not found: {input_file}")
            raise FileNotFoundError(f"Required input file not found: {input_file}")

        requests = cls._read_jsonl(input_file)

        bi = cls(
            model_name=model_id,
            job_name=job_name,
            region=region,
            bucket_name=bucket_name,
            role_arn=role_arn,
            session=session,
            output_dir=output_dir,
        )
        bi.job_arn = job_arn
        bi.requests = requests
        bi.job_status = response["status"]

        return bi

    except (KeyError, IndexError) as e:
        cls.logger.error(f"Invalid job response format: {e!s}")
        raise ValueError(f"Invalid job response format: {e!s}") from e
    except Exception as e:
        cls.logger.error(f"Failed to recover job details: {e!s}")
        raise RuntimeError(f"Failed to recover job details: {e!s}") from e

DeepSeekAdapter

Bases: OpenAICompatibleAdapter

Adapter for DeepSeek models in AWS Bedrock.

DeepSeek follows the OpenAI function-calling convention but differs in two ways:

• build_tool passes the raw Pydantic schema as parameters (no $defs inlining needed).
• prepare_model_input does not null anthropic_version or migrate the system prompt into the messages array.

Source code in src/llmbo/adapters/deepseek.py

class DeepSeekAdapter(OpenAICompatibleAdapter):
    """Adapter for DeepSeek models in AWS Bedrock.

    DeepSeek follows the OpenAI function-calling convention but differs
    in two ways:

    - ``build_tool`` passes the raw Pydantic schema as parameters
      (no ``$defs`` inlining needed).
    - ``prepare_model_input`` does not null ``anthropic_version`` or
      migrate the system prompt into the messages array.
    """

    logger = logging.getLogger(f"{__name__}.DeepSeekAdapter")
    _provider_name = "DeepSeek"

    @classmethod
    def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
        """Build a tool definition using the raw Pydantic schema.

        Args:
            output_model (type[BaseModel]): The Pydantic model to convert.

        Returns:
            dict[str, Any]: A tool definition dict.
        """
        cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

        return {
            "type": "function",
            "function": {
                "name": output_model.__name__,
                "description": (output_model.__doc__ or "Please fill in the schema"),
                "parameters": output_model.model_json_schema(),
            },
        }

    @classmethod
    def prepare_model_input(
        cls,
        model_input: ModelInput,
        output_model: type[BaseModel] | None = None,
    ) -> ModelInput:
        """Prepare model input for DeepSeek models.

        Unlike other OpenAI-compatible providers, DeepSeek does not
        require the system prompt to be moved into the messages array.

        Args:
            model_input (ModelInput): The original model input.
            output_model (type[BaseModel] | None): Optional Pydantic
                model defining the expected output structure.

        Returns:
            ModelInput: Modified model input.
        """
        cls.logger.debug(f"Preparing model input for {cls._provider_name}")

        if output_model:
            cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
            model_input.tools = [cls.build_tool(output_model)]
            model_input.tool_choice = "required"

        return model_input

build_tool(output_model) classmethod

Build a tool definition using the raw Pydantic schema.

Parameters:

Name	Type	Description	Default
output_model	type[BaseModel]	The Pydantic model to convert.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: A tool definition dict.

Source code in src/llmbo/adapters/deepseek.py

@classmethod
def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
    """Build a tool definition using the raw Pydantic schema.

    Args:
        output_model (type[BaseModel]): The Pydantic model to convert.

    Returns:
        dict[str, Any]: A tool definition dict.
    """
    cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

    return {
        "type": "function",
        "function": {
            "name": output_model.__name__,
            "description": (output_model.__doc__ or "Please fill in the schema"),
            "parameters": output_model.model_json_schema(),
        },
    }

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for DeepSeek models.

Unlike other OpenAI-compatible providers, DeepSeek does not require the system prompt to be moved into the messages array.

Parameters:

Name	Type	Description	Default
model_input	ModelInput	The original model input.	required
output_model	type[BaseModel] | None	Optional Pydantic model defining the expected output structure.	None

Returns:

Name	Type	Description
ModelInput	ModelInput	Modified model input.

Source code in src/llmbo/adapters/deepseek.py

@classmethod
def prepare_model_input(
    cls,
    model_input: ModelInput,
    output_model: type[BaseModel] | None = None,
) -> ModelInput:
    """Prepare model input for DeepSeek models.

    Unlike other OpenAI-compatible providers, DeepSeek does not
    require the system prompt to be moved into the messages array.

    Args:
        model_input (ModelInput): The original model input.
        output_model (type[BaseModel] | None): Optional Pydantic
            model defining the expected output structure.

    Returns:
        ModelInput: Modified model input.
    """
    cls.logger.debug(f"Preparing model input for {cls._provider_name}")

    if output_model:
        cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
        model_input.tools = [cls.build_tool(output_model)]
        model_input.tool_choice = "required"

    return model_input

LlamaAdapter

Bases: ModelProviderAdapter

Adapter for Meta Llama models (Llama 3 / 4) in AWS Bedrock.

This adapter handles: 1. Formatting the prompt using Meta's specific header tokens. 2. Enforcing JSON-only schema outputs. 3. Translating 'max_tokens' to Llama's native 'max_gen_len'.

Source code in src/llmbo/adapters/llama.py

class LlamaAdapter(ModelProviderAdapter):
    """Adapter for Meta Llama models (Llama 3 / 4) in AWS Bedrock.

    This adapter handles:
    1. Formatting the prompt using Meta's specific header tokens.
    2. Enforcing JSON-only schema outputs.
    3. Translating 'max_tokens' to Llama's native 'max_gen_len'.
    """

    logger = logging.getLogger(f"{__name__}.LlamaAdapter")

    @staticmethod
    def format_llama_prompt(user_prompt: str, system_prompt: str | None = None, tools: str | None = None) -> str:
        """Format a prompt using Meta's special header tokens.

        Assembles the ``<|begin_of_text|>``, system, user, and assistant
        header blocks into a single string that Llama's native endpoint
        expects.

        Args:
            user_prompt (str): The user's input or question.
            system_prompt (str | None): Optional system instructions.
            tools (str | None): Optional JSON schema string to inject
                as a structured-output constraint.

        Returns:
            str: The fully formatted prompt string.
        """

        prompt_parts = ["<|begin_of_text|>"]

        if system_prompt or tools:
            prompt_parts.append("<|start_header_id|>system<|end_header_id|>\n\n")
            if system_prompt:
                prompt_parts.append(f"{system_prompt}\n\n")
            if tools:
                prompt_parts.append(
                    "You must respond ONLY with a valid JSON object. "
                    "Do not include any conversational text, markdown formatting, or preamble. "
                    f"The JSON object must follow this exact schema:\n{tools}\n"
                )
            prompt_parts.append("<|eot_id|>")

        prompt_parts.append(f"<|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|>")
        prompt_parts.append("<|start_header_id|>assistant<|end_header_id|>\n\n")

        return "".join(prompt_parts)

    @staticmethod
    def _schema_to_string(output_model: type[BaseModel]) -> str:
        """Serialise a Pydantic model's JSON schema for prompt injection."""
        schema = output_model.model_json_schema()
        return json.dumps(schema, indent=2)

    @classmethod
    def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
        """Prepare model input for Meta Llama models.

        Converts the standard ``ModelInput`` into Llama's native format
        by building a single prompt string with special header tokens,
        translating ``max_tokens`` to ``max_gen_len``, and nullifying
        fields that Llama's endpoint does not accept.

        Args:
            model_input (ModelInput): The original model input.
            output_model (type[BaseModel] | None): Optional Pydantic
                model whose JSON schema is injected into the prompt
                to enforce structured output.

        Returns:
            ModelInput: Modified model input with Llama-specific fields
                populated and unsupported fields set to None.
        """

        cls.logger.debug("Preparing model input for Meta Llama")

        original_prompt = model_input.messages[0].get("content", "") if model_input.messages else ""
        tool = cls._schema_to_string(output_model) if output_model else None

        # Build the native prompt string
        system = model_input.system if isinstance(model_input.system, str) else None
        formatted_prompt = cls.format_llama_prompt(original_prompt, system, tool)

        # Inject the native Llama keys dynamically
        requested_tokens = model_input.max_tokens or 2048
        safe_max_gen_len = min(requested_tokens, 8192)

        model_input.prompt = formatted_prompt
        model_input.max_gen_len = safe_max_gen_len

        # Nullify the Anthropic/Converse keys so Bedrock doesn't reject them
        model_input.messages = None
        model_input.system = None
        model_input.max_tokens = None
        model_input.anthropic_version = None
        model_input.tools = None
        model_input.tool_choice = None

        return model_input

    @classmethod
    def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
        """Validate and parse output from Llama models.

        Llama's native endpoint returns free text in a ``generation``
        field. This method extracts the first JSON object found via
        regex and validates it against the provided Pydantic model.

        Args:
            result (dict[str, Any]): Raw model output from Llama.
            output_model (type[BaseModel]): Pydantic model to validate
                against.

        Returns:
            BaseModel | None: Validated model instance, or None if no
                JSON was found or validation fails.
        """

        cls.logger.debug(f"Validating result against {output_model.__name__} schema")

        # Bedrock's native Llama endpoint returns the text in a 'generation' key
        generation = result.get("generation", "")
        if not generation:
            cls.logger.debug("No 'generation' key found in result.")
            return None

        # Hunt for JSON brackets
        match = re.search(r"\{.*\}", generation, re.DOTALL)
        if not match:
            cls.logger.debug("Did not find anything that looked like JSON in the response")
            return None

        try:
            arguments = match.group(0)
            parsed_arguments = json.loads(arguments)
            return output_model(**parsed_arguments)
        except (json.JSONDecodeError, ValidationError) as e:
            cls.logger.debug(f"Validation failed: {e!s}")
            return None

_schema_to_string(output_model) staticmethod

Serialise a Pydantic model's JSON schema for prompt injection.

Source code in src/llmbo/adapters/llama.py

@staticmethod
def _schema_to_string(output_model: type[BaseModel]) -> str:
    """Serialise a Pydantic model's JSON schema for prompt injection."""
    schema = output_model.model_json_schema()
    return json.dumps(schema, indent=2)

format_llama_prompt(user_prompt, system_prompt=None, tools=None) staticmethod

Format a prompt using Meta's special header tokens.

Assembles the <|begin_of_text|>, system, user, and assistant header blocks into a single string that Llama's native endpoint expects.

Parameters:

Name	Type	Description	Default
user_prompt	str	The user's input or question.	required
system_prompt	str | None	Optional system instructions.	None
tools	str | None	Optional JSON schema string to inject as a structured-output constraint.	None

Returns:

Name	Type	Description
str	str	The fully formatted prompt string.

Source code in src/llmbo/adapters/llama.py

@staticmethod
def format_llama_prompt(user_prompt: str, system_prompt: str | None = None, tools: str | None = None) -> str:
    """Format a prompt using Meta's special header tokens.

    Assembles the ``<|begin_of_text|>``, system, user, and assistant
    header blocks into a single string that Llama's native endpoint
    expects.

    Args:
        user_prompt (str): The user's input or question.
        system_prompt (str | None): Optional system instructions.
        tools (str | None): Optional JSON schema string to inject
            as a structured-output constraint.

    Returns:
        str: The fully formatted prompt string.
    """

    prompt_parts = ["<|begin_of_text|>"]

    if system_prompt or tools:
        prompt_parts.append("<|start_header_id|>system<|end_header_id|>\n\n")
        if system_prompt:
            prompt_parts.append(f"{system_prompt}\n\n")
        if tools:
            prompt_parts.append(
                "You must respond ONLY with a valid JSON object. "
                "Do not include any conversational text, markdown formatting, or preamble. "
                f"The JSON object must follow this exact schema:\n{tools}\n"
            )
        prompt_parts.append("<|eot_id|>")

    prompt_parts.append(f"<|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|>")
    prompt_parts.append("<|start_header_id|>assistant<|end_header_id|>\n\n")

    return "".join(prompt_parts)

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for Meta Llama models.

Converts the standard ModelInput into Llama's native format by building a single prompt string with special header tokens, translating max_tokens to max_gen_len, and nullifying fields that Llama's endpoint does not accept.

Parameters:

Name	Type	Description	Default
model_input	ModelInput	The original model input.	required
output_model	type[BaseModel] | None	Optional Pydantic model whose JSON schema is injected into the prompt to enforce structured output.	None

Returns:

Name	Type	Description
ModelInput	ModelInput	Modified model input with Llama-specific fields populated and unsupported fields set to None.

Source code in src/llmbo/adapters/llama.py

@classmethod
def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
    """Prepare model input for Meta Llama models.

    Converts the standard ``ModelInput`` into Llama's native format
    by building a single prompt string with special header tokens,
    translating ``max_tokens`` to ``max_gen_len``, and nullifying
    fields that Llama's endpoint does not accept.

    Args:
        model_input (ModelInput): The original model input.
        output_model (type[BaseModel] | None): Optional Pydantic
            model whose JSON schema is injected into the prompt
            to enforce structured output.

    Returns:
        ModelInput: Modified model input with Llama-specific fields
            populated and unsupported fields set to None.
    """

    cls.logger.debug("Preparing model input for Meta Llama")

    original_prompt = model_input.messages[0].get("content", "") if model_input.messages else ""
    tool = cls._schema_to_string(output_model) if output_model else None

    # Build the native prompt string
    system = model_input.system if isinstance(model_input.system, str) else None
    formatted_prompt = cls.format_llama_prompt(original_prompt, system, tool)

    # Inject the native Llama keys dynamically
    requested_tokens = model_input.max_tokens or 2048
    safe_max_gen_len = min(requested_tokens, 8192)

    model_input.prompt = formatted_prompt
    model_input.max_gen_len = safe_max_gen_len

    # Nullify the Anthropic/Converse keys so Bedrock doesn't reject them
    model_input.messages = None
    model_input.system = None
    model_input.max_tokens = None
    model_input.anthropic_version = None
    model_input.tools = None
    model_input.tool_choice = None

    return model_input

validate_result(result, output_model) classmethod

Validate and parse output from Llama models.

Llama's native endpoint returns free text in a generation field. This method extracts the first JSON object found via regex and validates it against the provided Pydantic model.

Parameters:

Name	Type	Description	Default
result	dict[str, Any]	Raw model output from Llama.	required
output_model	type[BaseModel]	Pydantic model to validate against.	required

Returns:

Type	Description
BaseModel | None	BaseModel | None: Validated model instance, or None if no JSON was found or validation fails.

Source code in src/llmbo/adapters/llama.py

@classmethod
def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
    """Validate and parse output from Llama models.

    Llama's native endpoint returns free text in a ``generation``
    field. This method extracts the first JSON object found via
    regex and validates it against the provided Pydantic model.

    Args:
        result (dict[str, Any]): Raw model output from Llama.
        output_model (type[BaseModel]): Pydantic model to validate
            against.

    Returns:
        BaseModel | None: Validated model instance, or None if no
            JSON was found or validation fails.
    """

    cls.logger.debug(f"Validating result against {output_model.__name__} schema")

    # Bedrock's native Llama endpoint returns the text in a 'generation' key
    generation = result.get("generation", "")
    if not generation:
        cls.logger.debug("No 'generation' key found in result.")
        return None

    # Hunt for JSON brackets
    match = re.search(r"\{.*\}", generation, re.DOTALL)
    if not match:
        cls.logger.debug("Did not find anything that looked like JSON in the response")
        return None

    try:
        arguments = match.group(0)
        parsed_arguments = json.loads(arguments)
        return output_model(**parsed_arguments)
    except (json.JSONDecodeError, ValidationError) as e:
        cls.logger.debug(f"Validation failed: {e!s}")
        return None

Manifest

Bases: BaseModel

Job manifest details.

Uses extra="allow" so that new fields returned by the AWS Bedrock API (e.g. inputAudioSecond) are captured in model_extra instead of raising TypeError.

Source code in src/llmbo/models.py

class Manifest(BaseModel):
    """Job manifest details.

    Uses ``extra="allow"`` so that new fields returned by the AWS Bedrock API
    (e.g. ``inputAudioSecond``) are captured in ``model_extra`` instead of
    raising ``TypeError``.
    """

    model_config = ConfigDict(extra="allow")

    totalRecordCount: int
    processedRecordCount: int
    successRecordCount: int
    errorRecordCount: int
    inputTokenCount: int | None = None
    outputTokenCount: int | None = None

MistralAdapter

Bases: ModelProviderAdapter

Adapter for Mistral models in AWS Bedrock.

This adapter handles: 1. Formatting inputs for Mistral models 2. Building tool definitions in Mistral's format 3. Validating tool-use responses from Mistral models

Source code in src/llmbo/adapters/mistral.py

class MistralAdapter(ModelProviderAdapter):
    """Adapter for Mistral models in AWS Bedrock.

    This adapter handles:
    1. Formatting inputs for Mistral models
    2. Building tool definitions in Mistral's format
    3. Validating tool-use responses from Mistral models
    """

    logger = logging.getLogger(f"{__name__}.MistralAdapter")

    @staticmethod
    def format_mistral_prompt(user_prompt: str, system_prompt: str | None = None, tools: str | None = None) -> str:
        """
        Formats the user prompt, system prompt, and tool definitions for Mistral models.

        Parameters:
        - user_prompt (str): The user's input or question.
        - system_prompt (str, optional): The system's instructions or guidelines. Defaults to None.
        - tools (str, optional): Schema description. Defaults to None.

        Returns:
        - str: The formatted prompt ready for input into the Mistral model.
        """
        prompt_parts = ["<s>[INST]"]

        if tools:
            prompt_parts.append(
                "Reply with a JSON object. Reply only with the valid JSON object. "
                "The JSON object should follow the supplied schema. "
            )

        if system_prompt:
            prompt_parts.append(f"<<SYS>>\n{system_prompt}\n<</SYS>>\n")

        if tools:
            prompt_parts.append(f"{tools}\n")

        prompt_parts.append(f"{user_prompt} [/INST]")
        return "".join(prompt_parts)

    @classmethod
    def build_tool(cls, output_model: type[BaseModel]) -> str:  # type: ignore[override]
        """Build a tool definition in Mistral's format.

        Args:
            output_model: The Pydantic model to convert to a tool definition

        Returns:
            Dict with function definition for Mistral's tools format
        """
        cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

        schema = output_model.model_json_schema()

        tool = f"""The JSON Structure should be:\n\n\n{schema}\n\n\n"""

        cls.logger.debug(f"Created tool definition with name: {output_model.__name__}")
        return tool

    @classmethod
    def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
        """Prepare model input for Mistral models.

        Args:
            model_input: The original model input configuration
            output_model: The Pydantic model defining the expected output structure

        Returns:
            Modified model input with Mistral-specific configurations
        """
        cls.logger.debug("Preparing model input for Mistral")

        original_prompt = model_input.messages[0].get("content", "") if model_input.messages else ""
        if not original_prompt:
            cls.logger.debug("Didnt find any content to adapt")

        # Build tool from output_model and add it to model_input
        if output_model:
            cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
            tool = cls.build_tool(output_model)
        else:
            tool = None

        if model_input.messages:
            system = model_input.system if isinstance(model_input.system, str) else None
            model_input.messages[0]["content"] = cls.format_mistral_prompt(original_prompt, system, tool)

        # Mistral doesn't use anthropic_version, remove if set
        model_input.anthropic_version = None
        # It also doesnt support tools like this
        model_input.tools = None
        model_input.tool_choice = None
        # Or a system prompt
        model_input.system = None
        return model_input

    @classmethod
    def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
        """Validate and parse output from Mistral models.

        Extracts structured data from Mistral's tool-use response format and
        validates it against the provided Pydantic model.

        Args:
            result: Raw model output from Mistral
            output_model: Pydantic model to validate against

        Returns:
            Validated model instance or None if validation fails
        """
        cls.logger.debug(f"Validating result against {output_model.__name__} schema")

        # Check we have choices
        choices = result.get("choices", [])
        if not choices:
            cls.logger.debug("No expected 'choices' key in result.")
            return None

        # Check that we stopped on purpose
        if choices[0].get("finish_reason", "") != "stop":
            cls.logger.debug("Did not have 'stop' as the finish_reason.")
            return None

        # Check that the assistant returned a message
        if choices[0].get("message", {}).get("role", "") != "assistant":
            cls.logger.debug("Did not get the expected 'assistant' role.")
            return None

        # Check the content has something json looking.
        content = choices[0].get("message", {}).get("content", "")
        if content:
            match = re.search(r"\{.*\}", content, re.DOTALL)

        if not match:
            cls.logger.debug("Didnt find anything that looked like JSON in the response")
            return None

        try:
            arguments = match.group(0)
            parsed_arguments = json.loads(arguments)
        except json.JSONDecodeError:
            cls.logger.debug(f"Failed to parse function arguments as JSON: {arguments}")
            return None

        tool_name = parsed_arguments.get("title", "")
        if tool_name != output_model.__name__:
            cls.logger.debug(f"Wrong schema name in response, expected {output_model.__name__} got {tool_name}")

        try:
            validated_model = output_model(**parsed_arguments)
            cls.logger.debug("Validation successful.")
            return validated_model
        except ValidationError as e:
            cls.logger.debug(f"Validation failed: {e!s}")
            return None

build_tool(output_model) classmethod

Build a tool definition in Mistral's format.

Parameters:

Name	Type	Description	Default
output_model	type[BaseModel]	The Pydantic model to convert to a tool definition	required

Returns:

Type	Description
str	Dict with function definition for Mistral's tools format

Source code in src/llmbo/adapters/mistral.py

@classmethod
def build_tool(cls, output_model: type[BaseModel]) -> str:  # type: ignore[override]
    """Build a tool definition in Mistral's format.

    Args:
        output_model: The Pydantic model to convert to a tool definition

    Returns:
        Dict with function definition for Mistral's tools format
    """
    cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

    schema = output_model.model_json_schema()

    tool = f"""The JSON Structure should be:\n\n\n{schema}\n\n\n"""

    cls.logger.debug(f"Created tool definition with name: {output_model.__name__}")
    return tool

format_mistral_prompt(user_prompt, system_prompt=None, tools=None) staticmethod

Formats the user prompt, system prompt, and tool definitions for Mistral models.

Parameters: - user_prompt (str): The user's input or question. - system_prompt (str, optional): The system's instructions or guidelines. Defaults to None. - tools (str, optional): Schema description. Defaults to None.

Returns: - str: The formatted prompt ready for input into the Mistral model.

Source code in src/llmbo/adapters/mistral.py

@staticmethod
def format_mistral_prompt(user_prompt: str, system_prompt: str | None = None, tools: str | None = None) -> str:
    """
    Formats the user prompt, system prompt, and tool definitions for Mistral models.

    Parameters:
    - user_prompt (str): The user's input or question.
    - system_prompt (str, optional): The system's instructions or guidelines. Defaults to None.
    - tools (str, optional): Schema description. Defaults to None.

    Returns:
    - str: The formatted prompt ready for input into the Mistral model.
    """
    prompt_parts = ["<s>[INST]"]

    if tools:
        prompt_parts.append(
            "Reply with a JSON object. Reply only with the valid JSON object. "
            "The JSON object should follow the supplied schema. "
        )

    if system_prompt:
        prompt_parts.append(f"<<SYS>>\n{system_prompt}\n<</SYS>>\n")

    if tools:
        prompt_parts.append(f"{tools}\n")

    prompt_parts.append(f"{user_prompt} [/INST]")
    return "".join(prompt_parts)

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for Mistral models.

Parameters:

Name	Type	Description	Default
model_input	ModelInput	The original model input configuration	required
output_model	type[BaseModel] | None	The Pydantic model defining the expected output structure	None

Returns:

Type	Description
ModelInput	Modified model input with Mistral-specific configurations

Source code in src/llmbo/adapters/mistral.py

@classmethod
def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
    """Prepare model input for Mistral models.

    Args:
        model_input: The original model input configuration
        output_model: The Pydantic model defining the expected output structure

    Returns:
        Modified model input with Mistral-specific configurations
    """
    cls.logger.debug("Preparing model input for Mistral")

    original_prompt = model_input.messages[0].get("content", "") if model_input.messages else ""
    if not original_prompt:
        cls.logger.debug("Didnt find any content to adapt")

    # Build tool from output_model and add it to model_input
    if output_model:
        cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
        tool = cls.build_tool(output_model)
    else:
        tool = None

    if model_input.messages:
        system = model_input.system if isinstance(model_input.system, str) else None
        model_input.messages[0]["content"] = cls.format_mistral_prompt(original_prompt, system, tool)

    # Mistral doesn't use anthropic_version, remove if set
    model_input.anthropic_version = None
    # It also doesnt support tools like this
    model_input.tools = None
    model_input.tool_choice = None
    # Or a system prompt
    model_input.system = None
    return model_input

validate_result(result, output_model) classmethod

Validate and parse output from Mistral models.

Extracts structured data from Mistral's tool-use response format and validates it against the provided Pydantic model.

Parameters:

Name	Type	Description	Default
result	dict[str, Any]	Raw model output from Mistral	required
output_model	type[BaseModel]	Pydantic model to validate against	required

Returns:

Type	Description
BaseModel | None	Validated model instance or None if validation fails

Source code in src/llmbo/adapters/mistral.py

@classmethod
def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
    """Validate and parse output from Mistral models.

    Extracts structured data from Mistral's tool-use response format and
    validates it against the provided Pydantic model.

    Args:
        result: Raw model output from Mistral
        output_model: Pydantic model to validate against

    Returns:
        Validated model instance or None if validation fails
    """
    cls.logger.debug(f"Validating result against {output_model.__name__} schema")

    # Check we have choices
    choices = result.get("choices", [])
    if not choices:
        cls.logger.debug("No expected 'choices' key in result.")
        return None

    # Check that we stopped on purpose
    if choices[0].get("finish_reason", "") != "stop":
        cls.logger.debug("Did not have 'stop' as the finish_reason.")
        return None

    # Check that the assistant returned a message
    if choices[0].get("message", {}).get("role", "") != "assistant":
        cls.logger.debug("Did not get the expected 'assistant' role.")
        return None

    # Check the content has something json looking.
    content = choices[0].get("message", {}).get("content", "")
    if content:
        match = re.search(r"\{.*\}", content, re.DOTALL)

    if not match:
        cls.logger.debug("Didnt find anything that looked like JSON in the response")
        return None

    try:
        arguments = match.group(0)
        parsed_arguments = json.loads(arguments)
    except json.JSONDecodeError:
        cls.logger.debug(f"Failed to parse function arguments as JSON: {arguments}")
        return None

    tool_name = parsed_arguments.get("title", "")
    if tool_name != output_model.__name__:
        cls.logger.debug(f"Wrong schema name in response, expected {output_model.__name__} got {tool_name}")

    try:
        validated_model = output_model(**parsed_arguments)
        cls.logger.debug("Validation successful.")
        return validated_model
    except ValidationError as e:
        cls.logger.debug(f"Validation failed: {e!s}")
        return None

MistralFunctionAdapter

Bases: ModelProviderAdapter

Adapter for Mistral models using function calling in AWS Bedrock.

This adapter handles: 1. Building tool definitions in Mistral's function-calling format. 2. Migrating the system prompt into the messages array. 3. Enforcing Mistral's 8192 token ceiling. 4. Validating tool-use responses with granular debug logging.

Source code in src/llmbo/adapters/mistral_function_calling.py

class MistralFunctionAdapter(ModelProviderAdapter):
    """Adapter for Mistral models using function calling in AWS Bedrock.

    This adapter handles:
    1. Building tool definitions in Mistral's function-calling format.
    2. Migrating the system prompt into the messages array.
    3. Enforcing Mistral's 8192 token ceiling.
    4. Validating tool-use responses with granular debug logging.
    """

    logger = logging.getLogger(f"{__name__}.MistralFunctionAdapter")

    @classmethod
    def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
        """Build a tool definition in Mistral's function-calling format.

        Preserves any ``$defs`` block generated by Pydantic, as Mistral
        can resolve nested schema references.

        Args:
            output_model (type[BaseModel]): The Pydantic model to convert
                to a tool definition.

        Returns:
            dict[str, Any]: Tool definition dict with function name,
                description, and parameters.
        """

        cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

        schema = output_model.model_json_schema()
        parameters = {
            "type": "object",
            "properties": schema.get("properties", {}),
            "required": schema.get("required", []),
        }

        # Keep nested definitions if Pydantic generated them
        if "$defs" in schema:
            parameters["$defs"] = schema["$defs"]

        return {
            "type": "function",
            "function": {
                "name": output_model.__name__,
                "description": schema.get("description", ""),
                "parameters": parameters,
            },
        }

    @classmethod
    def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
        """Prepare model input for Mistral function-calling models.

        Moves the system prompt into the messages array, caps
        ``max_tokens`` at 8192, and attaches a tool definition when
        an output model is provided.

        Args:
            model_input (ModelInput): The original model input.
            output_model (type[BaseModel] | None): Optional Pydantic
                model defining the expected output structure.

        Returns:
            ModelInput: Modified model input with Mistral-specific
                configurations applied.
        """

        cls.logger.debug("Preparing model input for Mistral")

        model_input.anthropic_version = None

        # --- THE SYSTEM PROMPT FIX ---
        # Mistral strictly expects the system prompt inside the messages array
        if model_input.system and isinstance(model_input.system, str):
            if not model_input.messages:
                model_input.messages = []

            # Prepend the system prompt to the messages list
            model_input.messages.insert(0, {"role": "system", "content": model_input.system})

            # Delete the top-level key so AWS doesn't throw the 'extra_forbidden' error
            model_input.system = None
        # -----------------------------

        # Enforce Mistral's token ceiling
        if model_input.max_tokens and model_input.max_tokens > 8192:
            model_input.max_tokens = 8192

        if output_model:
            cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
            model_input.tools = [cls.build_tool(output_model)]
            model_input.tool_choice = "any"

        return model_input

    @classmethod
    def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
        """Validate and parse output from Mistral function-calling models.

        Performs granular checks on the response structure — finish
        reason, assistant role, tool call count, and tool name — with
        debug logging at each step to aid batch failure diagnosis.

        Args:
            result (dict[str, Any]): Raw model output from Mistral.
            output_model (type[BaseModel]): Pydantic model to validate
                against.

        Returns:
            BaseModel | None: Validated model instance, or None if any
                check fails.
        """

        cls.logger.debug(f"Validating result against {output_model.__name__} schema")

        # Check we have choices
        choices = result.get("choices", [])
        if not choices:
            cls.logger.debug("No expected 'choices' key in result.")
            return None

        choice = choices[0]

        # Check that we stopped on purpose
        finish_reason = choice.get("finish_reason", "")
        if finish_reason != "tool_calls":
            cls.logger.debug(f"Expected 'tool_calls' as finish_reason, got '{finish_reason}'.")
            return None

        # Check that the assistant returned a message
        message = choice.get("message", {})
        if message.get("role", "") != "assistant":
            cls.logger.debug("Did not get the expected 'assistant' role.")
            return None

        # Check that the assistant returned tool calls
        tools = message.get("tool_calls", [])
        if not tools:
            cls.logger.debug("No tool calls found in assistant message.")
            return None

        if len(tools) != 1:
            cls.logger.debug(f"Expected exactly 1 tool call, got {len(tools)}.")
            return None

        # Check tool name matches expected model
        function = tools[0].get("function", {})
        tool_name = function.get("name", "")
        if tool_name != output_model.__name__:
            cls.logger.debug(f"Wrong tool name in response, expected '{output_model.__name__}' got '{tool_name}'.")
            return None

        # Parse and validate arguments
        try:
            arguments = function.get("arguments", "{}")
            parsed_arguments = json.loads(arguments)
        except json.JSONDecodeError:
            cls.logger.debug(f"Failed to parse function arguments as JSON: {arguments}")
            return None

        try:
            validated_model = output_model(**parsed_arguments)
            cls.logger.debug("Validation successful.")
            return validated_model
        except ValidationError as e:
            cls.logger.debug(f"Validation failed: {e!s}")
            return None

build_tool(output_model) classmethod

Build a tool definition in Mistral's function-calling format.

Preserves any $defs block generated by Pydantic, as Mistral can resolve nested schema references.

Parameters:

Name	Type	Description	Default
output_model	type[BaseModel]	The Pydantic model to convert to a tool definition.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: Tool definition dict with function name, description, and parameters.

Source code in src/llmbo/adapters/mistral_function_calling.py

@classmethod
def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
    """Build a tool definition in Mistral's function-calling format.

    Preserves any ``$defs`` block generated by Pydantic, as Mistral
    can resolve nested schema references.

    Args:
        output_model (type[BaseModel]): The Pydantic model to convert
            to a tool definition.

    Returns:
        dict[str, Any]: Tool definition dict with function name,
            description, and parameters.
    """

    cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

    schema = output_model.model_json_schema()
    parameters = {
        "type": "object",
        "properties": schema.get("properties", {}),
        "required": schema.get("required", []),
    }

    # Keep nested definitions if Pydantic generated them
    if "$defs" in schema:
        parameters["$defs"] = schema["$defs"]

    return {
        "type": "function",
        "function": {
            "name": output_model.__name__,
            "description": schema.get("description", ""),
            "parameters": parameters,
        },
    }

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for Mistral function-calling models.

Moves the system prompt into the messages array, caps max_tokens at 8192, and attaches a tool definition when an output model is provided.

Parameters:

Name	Type	Description	Default
model_input	ModelInput	The original model input.	required
output_model	type[BaseModel] | None	Optional Pydantic model defining the expected output structure.	None

Returns:

Name	Type	Description
ModelInput	ModelInput	Modified model input with Mistral-specific configurations applied.

Source code in src/llmbo/adapters/mistral_function_calling.py

@classmethod
def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
    """Prepare model input for Mistral function-calling models.

    Moves the system prompt into the messages array, caps
    ``max_tokens`` at 8192, and attaches a tool definition when
    an output model is provided.

    Args:
        model_input (ModelInput): The original model input.
        output_model (type[BaseModel] | None): Optional Pydantic
            model defining the expected output structure.

    Returns:
        ModelInput: Modified model input with Mistral-specific
            configurations applied.
    """

    cls.logger.debug("Preparing model input for Mistral")

    model_input.anthropic_version = None

    # --- THE SYSTEM PROMPT FIX ---
    # Mistral strictly expects the system prompt inside the messages array
    if model_input.system and isinstance(model_input.system, str):
        if not model_input.messages:
            model_input.messages = []

        # Prepend the system prompt to the messages list
        model_input.messages.insert(0, {"role": "system", "content": model_input.system})

        # Delete the top-level key so AWS doesn't throw the 'extra_forbidden' error
        model_input.system = None
    # -----------------------------

    # Enforce Mistral's token ceiling
    if model_input.max_tokens and model_input.max_tokens > 8192:
        model_input.max_tokens = 8192

    if output_model:
        cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
        model_input.tools = [cls.build_tool(output_model)]
        model_input.tool_choice = "any"

    return model_input

validate_result(result, output_model) classmethod

Validate and parse output from Mistral function-calling models.

Performs granular checks on the response structure — finish reason, assistant role, tool call count, and tool name — with debug logging at each step to aid batch failure diagnosis.

Parameters:

Name	Type	Description	Default
result	dict[str, Any]	Raw model output from Mistral.	required
output_model	type[BaseModel]	Pydantic model to validate against.	required

Returns:

Type	Description
BaseModel | None	BaseModel | None: Validated model instance, or None if any check fails.

Source code in src/llmbo/adapters/mistral_function_calling.py

@classmethod
def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
    """Validate and parse output from Mistral function-calling models.

    Performs granular checks on the response structure — finish
    reason, assistant role, tool call count, and tool name — with
    debug logging at each step to aid batch failure diagnosis.

    Args:
        result (dict[str, Any]): Raw model output from Mistral.
        output_model (type[BaseModel]): Pydantic model to validate
            against.

    Returns:
        BaseModel | None: Validated model instance, or None if any
            check fails.
    """

    cls.logger.debug(f"Validating result against {output_model.__name__} schema")

    # Check we have choices
    choices = result.get("choices", [])
    if not choices:
        cls.logger.debug("No expected 'choices' key in result.")
        return None

    choice = choices[0]

    # Check that we stopped on purpose
    finish_reason = choice.get("finish_reason", "")
    if finish_reason != "tool_calls":
        cls.logger.debug(f"Expected 'tool_calls' as finish_reason, got '{finish_reason}'.")
        return None

    # Check that the assistant returned a message
    message = choice.get("message", {})
    if message.get("role", "") != "assistant":
        cls.logger.debug("Did not get the expected 'assistant' role.")
        return None

    # Check that the assistant returned tool calls
    tools = message.get("tool_calls", [])
    if not tools:
        cls.logger.debug("No tool calls found in assistant message.")
        return None

    if len(tools) != 1:
        cls.logger.debug(f"Expected exactly 1 tool call, got {len(tools)}.")
        return None

    # Check tool name matches expected model
    function = tools[0].get("function", {})
    tool_name = function.get("name", "")
    if tool_name != output_model.__name__:
        cls.logger.debug(f"Wrong tool name in response, expected '{output_model.__name__}' got '{tool_name}'.")
        return None

    # Parse and validate arguments
    try:
        arguments = function.get("arguments", "{}")
        parsed_arguments = json.loads(arguments)
    except json.JSONDecodeError:
        cls.logger.debug(f"Failed to parse function arguments as JSON: {arguments}")
        return None

    try:
        validated_model = output_model(**parsed_arguments)
        cls.logger.debug("Validation successful.")
        return validated_model
    except ValidationError as e:
        cls.logger.debug(f"Validation failed: {e!s}")
        return None

ModelAdapterRegistry

Registry for model provider adapters.

This registry maps model name patterns to their corresponding adapter classes. Users can register custom adapters for new model providers or to override existing implementations.

Example

Register a custom adapter for a new model

ModelAdapterRegistry.register("my-custom-model", MyCustomAdapter)

Source code in src/llmbo/registry.py

class ModelAdapterRegistry:
    """Registry for model provider adapters.

    This registry maps model name patterns to their corresponding adapter classes.
    Users can register custom adapters for new model providers or to override
    existing implementations.

    Example:
        >>> # Register a custom adapter for a new model
        >>> ModelAdapterRegistry.register("my-custom-model", MyCustomAdapter)
    """

    _adapters: ClassVar[list[tuple[Pattern, type[ModelProviderAdapter]]]] = []
    logger = logging.getLogger(__name__)

    @classmethod
    def register(cls, pattern: str, adapter_class: type[ModelProviderAdapter]) -> None:
        """Register an adapter class for a specific model pattern.

        Args:
            pattern: Regex pattern to match against model names
            adapter_class: The adapter class to use for matching models

        Raises:
            TypeError: If adapter_class is not a subclass of ModelProviderAdapter
        """
        # Add type validation to ensure adapter_class is a proper ModelProviderAdapter
        if not issubclass(adapter_class, ModelProviderAdapter):
            cls.logger.error(f"Adapter class must be a subclass of ModelProviderAdapter, got {adapter_class.__name__}")
            raise TypeError(f"Adapter class must be a subclass of ModelProviderAdapter, got {adapter_class.__name__}")

        compiled_pattern = re.compile(pattern)

        # Check for duplicate pattern and log a warning
        for i, (existing_pattern, _) in enumerate(cls._adapters):
            if existing_pattern.pattern == compiled_pattern.pattern:
                cls.logger.warning(f"Adapter for pattern '{pattern}' is being replaced with {adapter_class.__name__}")
                # Remove the existing adapter with the same pattern
                cls._adapters.pop(i)
                break

        # Add new adapter to the beginning of the list for higher precedence
        cls._adapters.insert(0, (compiled_pattern, adapter_class))
        cls.logger.info(f"Registered adapter {adapter_class.__name__} for pattern '{pattern}'")

    @classmethod
    def get_adapter(cls, model_name: str) -> type[ModelProviderAdapter]:
        """Get the appropriate adapter for a model name.

        Args:
            model_name: The model name/ID to find an adapter for

        Returns:
            An adapter class for the given model, or the default adapter if no pattern
            is found
        """
        for pattern, adapter in cls._adapters:
            if pattern.search(model_name):
                return adapter

        cls.logger.warning(
            f"No pattern found for {model_name}, returning default ModelAdapter. "
            "This model is unsupported it may not work as expected.",
        )
        return DefaultAdapter

get_adapter(model_name) classmethod

Get the appropriate adapter for a model name.

Parameters:

Name	Type	Description	Default
model_name	str	The model name/ID to find an adapter for	required

Returns:

Type	Description
type[ModelProviderAdapter]	An adapter class for the given model, or the default adapter if no pattern
type[ModelProviderAdapter]	is found

Source code in src/llmbo/registry.py

@classmethod
def get_adapter(cls, model_name: str) -> type[ModelProviderAdapter]:
    """Get the appropriate adapter for a model name.

    Args:
        model_name: The model name/ID to find an adapter for

    Returns:
        An adapter class for the given model, or the default adapter if no pattern
        is found
    """
    for pattern, adapter in cls._adapters:
        if pattern.search(model_name):
            return adapter

    cls.logger.warning(
        f"No pattern found for {model_name}, returning default ModelAdapter. "
        "This model is unsupported it may not work as expected.",
    )
    return DefaultAdapter

register(pattern, adapter_class) classmethod

Register an adapter class for a specific model pattern.

Parameters:

Name	Type	Description	Default
pattern	str	Regex pattern to match against model names	required
adapter_class	type[ModelProviderAdapter]	The adapter class to use for matching models	required

Raises:

Type	Description
TypeError	If adapter_class is not a subclass of ModelProviderAdapter

Source code in src/llmbo/registry.py

@classmethod
def register(cls, pattern: str, adapter_class: type[ModelProviderAdapter]) -> None:
    """Register an adapter class for a specific model pattern.

    Args:
        pattern: Regex pattern to match against model names
        adapter_class: The adapter class to use for matching models

    Raises:
        TypeError: If adapter_class is not a subclass of ModelProviderAdapter
    """
    # Add type validation to ensure adapter_class is a proper ModelProviderAdapter
    if not issubclass(adapter_class, ModelProviderAdapter):
        cls.logger.error(f"Adapter class must be a subclass of ModelProviderAdapter, got {adapter_class.__name__}")
        raise TypeError(f"Adapter class must be a subclass of ModelProviderAdapter, got {adapter_class.__name__}")

    compiled_pattern = re.compile(pattern)

    # Check for duplicate pattern and log a warning
    for i, (existing_pattern, _) in enumerate(cls._adapters):
        if existing_pattern.pattern == compiled_pattern.pattern:
            cls.logger.warning(f"Adapter for pattern '{pattern}' is being replaced with {adapter_class.__name__}")
            # Remove the existing adapter with the same pattern
            cls._adapters.pop(i)
            break

    # Add new adapter to the beginning of the list for higher precedence
    cls._adapters.insert(0, (compiled_pattern, adapter_class))
    cls.logger.info(f"Registered adapter {adapter_class.__name__} for pattern '{pattern}'")

ModelInput dataclass

Configuration class for AWS Bedrock model inputs.

This class defines the structure and parameters for model invocation requests following AWS Bedrock's expected format. Provider-specific adapters may reshape or nullify fields as needed for their API.

See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html

Attributes:

Name	Type	Description
messages	list[dict] | None	List of message objects with role and content. Defaults to None.
anthropic_version	str | None	Version string for Anthropic models. Defaults to "bedrock-2023-05-31".
max_tokens	int | None	Maximum number of tokens in the response. Defaults to 2000.
system	str | list[dict[str, Any]] | None	System message for the model. A string for most providers; reshaped to a list of content blocks for Converse API models (e.g. Nova).
stop_sequences	list[str] | None	Custom stop sequences.
temperature	float | None	Sampling temperature.
top_p	float | None	Nucleus sampling parameter.
top_k	int | None	Top-k sampling parameter.
tools	list[dict] | None	Tool definitions for structured outputs.
tool_choice	ToolChoice | str | None	Tool selection configuration.
prompt	str | None	Native text prompt for models that use a single string instead of a messages array (e.g. Llama).
max_gen_len	int | None	Maximum generation length for models that use this parameter instead of max_tokens (e.g. Llama).
inferenceConfig	dict[str, Any] | None	Inference configuration for Converse API models (e.g. Nova).
toolConfig	dict[str, Any] | None	Tool configuration for Converse API models (e.g. Nova).
additionalModelRequestFields	dict[str, Any] | None	Extra provider request fields for Converse API models; used for parameters like topK that live outside inferenceConfig (e.g. Nova).

Source code in src/llmbo/models.py

@dataclass
class ModelInput:
    """Configuration class for AWS Bedrock model inputs.

    This class defines the structure and parameters for model invocation requests
    following AWS Bedrock's expected format. Provider-specific adapters may
    reshape or nullify fields as needed for their API.

    See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html

    Attributes:
        messages (list[dict] | None): List of message objects with role
            and content. Defaults to None.
        anthropic_version (str | None): Version string for Anthropic
            models. Defaults to "bedrock-2023-05-31".
        max_tokens (int | None): Maximum number of tokens in the
            response. Defaults to 2000.
        system (str | list[dict[str, Any]] | None): System message for
            the model. A string for most providers; reshaped to a list
            of content blocks for Converse API models (e.g. Nova).
        stop_sequences (list[str] | None): Custom stop sequences.
        temperature (float | None): Sampling temperature.
        top_p (float | None): Nucleus sampling parameter.
        top_k (int | None): Top-k sampling parameter.
        tools (list[dict] | None): Tool definitions for structured
            outputs.
        tool_choice (ToolChoice | str | None): Tool selection
            configuration.
        prompt (str | None): Native text prompt for models that use a
            single string instead of a messages array (e.g. Llama).
        max_gen_len (int | None): Maximum generation length for models
            that use this parameter instead of max_tokens (e.g. Llama).
        inferenceConfig (dict[str, Any] | None): Inference
            configuration for Converse API models (e.g. Nova).
        toolConfig (dict[str, Any] | None): Tool configuration for
            Converse API models (e.g. Nova).
        additionalModelRequestFields (dict[str, Any] | None): Extra
            provider request fields for Converse API models; used for
            parameters like topK that live outside inferenceConfig
            (e.g. Nova).
    """

    # These are required
    messages: list[dict] | None = None
    anthropic_version: str | None = "bedrock-2023-05-31"
    max_tokens: int | None = 2000

    system: str | list[dict[str, Any]] | None = None
    stop_sequences: list[str] | None = None
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None

    tools: list[dict] | None = None
    tool_choice: ToolChoice | str | None = None

    # Provider-specific fields (used by adapters that diverge from the
    # Anthropic / Messages API shape, e.g. Llama, Nova)
    prompt: str | None = None
    max_gen_len: int | None = None
    inferenceConfig: dict[str, Any] | None = None
    toolConfig: dict[str, Any] | None = None
    additionalModelRequestFields: dict[str, Any] | None = None

    def to_dict(self):
        """Convert to dict."""
        result = {k: v for k, v in self.__dict__.items() if v is not None}
        if isinstance(self.tool_choice, ToolChoice):
            result["tool_choice"] = self.tool_choice.__dict__
        return result

    def to_json(self):
        """Convert to json string."""
        return json.dumps(self.to_dict())

to_dict()

Convert to dict.

Source code in src/llmbo/models.py

def to_dict(self):
    """Convert to dict."""
    result = {k: v for k, v in self.__dict__.items() if v is not None}
    if isinstance(self.tool_choice, ToolChoice):
        result["tool_choice"] = self.tool_choice.__dict__
    return result

to_json()

Convert to json string.

Source code in src/llmbo/models.py

def to_json(self):
    """Convert to json string."""
    return json.dumps(self.to_dict())

NovaAdapter

Bases: ModelProviderAdapter

Adapter for Amazon Nova models using the Bedrock Converse API.

This adapter handles: 1. Building tool definitions in the Converse toolSpec format. 2. Reshaping messages, system prompts, and token limits into Converse API JSON arrays. 3. Inlining $defs to resolve Pydantic $ref pointers that Nova cannot follow. 4. Validating toolUse responses from the Converse API.

Source code in src/llmbo/adapters/nova.py

class NovaAdapter(ModelProviderAdapter):
    """Adapter for Amazon Nova models using the Bedrock Converse API.

    This adapter handles:
    1. Building tool definitions in the Converse ``toolSpec`` format.
    2. Reshaping messages, system prompts, and token limits into
       Converse API JSON arrays.
    3. Inlining ``$defs`` to resolve Pydantic ``$ref`` pointers that
       Nova cannot follow.
    4. Validating ``toolUse`` responses from the Converse API.
    """

    logger = logging.getLogger(f"{__name__}.NovaAdapter")

    @classmethod
    def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
        """Build a tool definition in the Converse API ``toolSpec`` format.

        Inlines any ``$defs`` generated by Pydantic and provides a
        fallback description if the model's schema has none, since the
        Converse API rejects empty description strings.

        Args:
            output_model (type[BaseModel]): The Pydantic model to convert
                to a tool definition.

        Returns:
            dict[str, Any]: A Converse API tool definition with
                ``toolSpec``, ``inputSchema``, and resolved properties.
        """

        cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

        schema = output_model.model_json_schema()
        properties = schema.get("properties", {})

        # Inline the $defs to prevent the "lazy pointer" issue we saw with OSS models
        if "$defs" in schema:
            for _, prop_val in properties.items():
                if prop_val.get("type") == "array" and "$ref" in prop_val.get("items", {}):
                    ref_name = prop_val["items"]["$ref"].split("/")[-1]
                    if ref_name in schema["$defs"]:
                        prop_val["items"] = schema["$defs"][ref_name]

        parameters = {
            "type": "object",
            "properties": properties,
            "required": schema.get("required", []),
        }

        # Converse API strictly rejects zero-length strings for descriptions.
        description = schema.get("description", "").strip()
        if not description:
            description = f"Extract structured data matching the {output_model.__name__} schema."

        # Converse API wraps tools inside a "toolSpec" and "inputSchema -> json"
        return {
            "toolSpec": {"name": output_model.__name__, "description": description, "inputSchema": {"json": parameters}}
        }

    @classmethod
    def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
        """Prepare model input for Amazon Nova via the Converse API.

        Reshapes the standard ``ModelInput`` into the Converse API
        format: messages become content-block arrays, the system prompt
        becomes a text array, and sampling parameters (max_tokens,
        temperature, top_p) move into inferenceConfig. Nova does not
        support top_k, so it is dropped.

        Args:
            model_input (ModelInput): The original model input.
            output_model (type[BaseModel] | None): Optional Pydantic
                model defining the expected output structure.

        Returns:
            ModelInput: Modified model input with Converse API fields
                populated and legacy fields nullified.
        """

        cls.logger.debug("Preparing model input for Amazon Nova")

        # Clean up legacy artifacts
        model_input.anthropic_version = None

        # 1. Reshape the User Messages into Converse JSONArrays
        if model_input.messages:
            for msg in model_input.messages:
                if isinstance(msg.get("content"), str):
                    msg["content"] = [{"text": msg["content"]}]

        # 2. Reshape the System Prompt into a Converse JSONArray
        if model_input.system and isinstance(model_input.system, str):
            model_input.system = [{"text": model_input.system}]

        # 3. Reshape sampling params into inferenceConfig
        #    The Converse API rejects top-level temperature/top_p, so they
        #    must be nested. top_k is unsupported by Nova and is dropped.
        inference_config: dict[str, Any] = {}
        if model_input.max_tokens:
            inference_config["maxTokens"] = model_input.max_tokens
        if model_input.temperature is not None:
            inference_config["temperature"] = model_input.temperature
        if model_input.top_p is not None:
            inference_config["topP"] = model_input.top_p

        if inference_config:
            model_input.inferenceConfig = inference_config

        # Remove the top-level keys so to_dict() does not emit them.
        # top_k is nulled (not routed) as Nova's Converse API rejects it.
        model_input.max_tokens = None
        model_input.temperature = None
        model_input.top_p = None
        model_input.top_k = None

        # 4. Reshape tools into toolConfig
        if output_model:
            cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
            tool_config = {
                "tools": [cls.build_tool(output_model)],
                "toolChoice": {"tool": {"name": output_model.__name__}},
            }
            # Dynamically attach the toolConfig and delete the legacy keys
            model_input.toolConfig = tool_config
            model_input.tools = None
            model_input.tool_choice = None

        return model_input

    @classmethod
    def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
        """Validate and parse output from Amazon Nova's Converse API.

        Searches the response content blocks for a ``toolUse`` entry
        and validates its input against the provided Pydantic model.

        Args:
            result (dict[str, Any]): Raw model output from the Converse
                API.
            output_model (type[BaseModel]): Pydantic model to validate
                against.

        Returns:
            BaseModel | None: Validated model instance, or None if no
                ``toolUse`` block was found or validation fails.
        """

        cls.logger.debug(f"Validating result against {output_model.__name__} schema")

        # Converse API returns data in: output -> message -> content -> toolUse
        output = result.get("output", {})
        message = output.get("message", {})
        content_blocks = message.get("content", [])

        for block in content_blocks:
            if "toolUse" in block:
                try:
                    return output_model(**block["toolUse"]["input"])
                except ValidationError as e:
                    cls.logger.debug(f"Validation failed: {e!s}")
                    return None

        cls.logger.debug("Result contains no toolUse. The model hallucinated text.")
        return None

build_tool(output_model) classmethod

Build a tool definition in the Converse API toolSpec format.

Inlines any $defs generated by Pydantic and provides a fallback description if the model's schema has none, since the Converse API rejects empty description strings.

Parameters:

Name	Type	Description	Default
output_model	type[BaseModel]	The Pydantic model to convert to a tool definition.	required

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: A Converse API tool definition with toolSpec, inputSchema, and resolved properties.

Source code in src/llmbo/adapters/nova.py

@classmethod
def build_tool(cls, output_model: type[BaseModel]) -> dict[str, Any]:
    """Build a tool definition in the Converse API ``toolSpec`` format.

    Inlines any ``$defs`` generated by Pydantic and provides a
    fallback description if the model's schema has none, since the
    Converse API rejects empty description strings.

    Args:
        output_model (type[BaseModel]): The Pydantic model to convert
            to a tool definition.

    Returns:
        dict[str, Any]: A Converse API tool definition with
            ``toolSpec``, ``inputSchema``, and resolved properties.
    """

    cls.logger.debug(f"Building tool definition for model: {output_model.__name__}")

    schema = output_model.model_json_schema()
    properties = schema.get("properties", {})

    # Inline the $defs to prevent the "lazy pointer" issue we saw with OSS models
    if "$defs" in schema:
        for _, prop_val in properties.items():
            if prop_val.get("type") == "array" and "$ref" in prop_val.get("items", {}):
                ref_name = prop_val["items"]["$ref"].split("/")[-1]
                if ref_name in schema["$defs"]:
                    prop_val["items"] = schema["$defs"][ref_name]

    parameters = {
        "type": "object",
        "properties": properties,
        "required": schema.get("required", []),
    }

    # Converse API strictly rejects zero-length strings for descriptions.
    description = schema.get("description", "").strip()
    if not description:
        description = f"Extract structured data matching the {output_model.__name__} schema."

    # Converse API wraps tools inside a "toolSpec" and "inputSchema -> json"
    return {
        "toolSpec": {"name": output_model.__name__, "description": description, "inputSchema": {"json": parameters}}
    }

prepare_model_input(model_input, output_model=None) classmethod

Prepare model input for Amazon Nova via the Converse API.

Reshapes the standard ModelInput into the Converse API format: messages become content-block arrays, the system prompt becomes a text array, and sampling parameters (max_tokens, temperature, top_p) move into inferenceConfig. Nova does not support top_k, so it is dropped.

Parameters:

Name	Type	Description	Default
model_input	ModelInput	The original model input.	required
output_model	type[BaseModel] | None	Optional Pydantic model defining the expected output structure.	None

Returns:

Name	Type	Description
ModelInput	ModelInput	Modified model input with Converse API fields populated and legacy fields nullified.

Source code in src/llmbo/adapters/nova.py

@classmethod
def prepare_model_input(cls, model_input: ModelInput, output_model: type[BaseModel] | None = None) -> ModelInput:
    """Prepare model input for Amazon Nova via the Converse API.

    Reshapes the standard ``ModelInput`` into the Converse API
    format: messages become content-block arrays, the system prompt
    becomes a text array, and sampling parameters (max_tokens,
    temperature, top_p) move into inferenceConfig. Nova does not
    support top_k, so it is dropped.

    Args:
        model_input (ModelInput): The original model input.
        output_model (type[BaseModel] | None): Optional Pydantic
            model defining the expected output structure.

    Returns:
        ModelInput: Modified model input with Converse API fields
            populated and legacy fields nullified.
    """

    cls.logger.debug("Preparing model input for Amazon Nova")

    # Clean up legacy artifacts
    model_input.anthropic_version = None

    # 1. Reshape the User Messages into Converse JSONArrays
    if model_input.messages:
        for msg in model_input.messages:
            if isinstance(msg.get("content"), str):
                msg["content"] = [{"text": msg["content"]}]

    # 2. Reshape the System Prompt into a Converse JSONArray
    if model_input.system and isinstance(model_input.system, str):
        model_input.system = [{"text": model_input.system}]

    # 3. Reshape sampling params into inferenceConfig
    #    The Converse API rejects top-level temperature/top_p, so they
    #    must be nested. top_k is unsupported by Nova and is dropped.
    inference_config: dict[str, Any] = {}
    if model_input.max_tokens:
        inference_config["maxTokens"] = model_input.max_tokens
    if model_input.temperature is not None:
        inference_config["temperature"] = model_input.temperature
    if model_input.top_p is not None:
        inference_config["topP"] = model_input.top_p

    if inference_config:
        model_input.inferenceConfig = inference_config

    # Remove the top-level keys so to_dict() does not emit them.
    # top_k is nulled (not routed) as Nova's Converse API rejects it.
    model_input.max_tokens = None
    model_input.temperature = None
    model_input.top_p = None
    model_input.top_k = None

    # 4. Reshape tools into toolConfig
    if output_model:
        cls.logger.debug(f"Adding tool definition for {output_model.__name__}")
        tool_config = {
            "tools": [cls.build_tool(output_model)],
            "toolChoice": {"tool": {"name": output_model.__name__}},
        }
        # Dynamically attach the toolConfig and delete the legacy keys
        model_input.toolConfig = tool_config
        model_input.tools = None
        model_input.tool_choice = None

    return model_input

validate_result(result, output_model) classmethod

Validate and parse output from Amazon Nova's Converse API.

Searches the response content blocks for a toolUse entry and validates its input against the provided Pydantic model.

Parameters:

Name	Type	Description	Default
result	dict[str, Any]	Raw model output from the Converse API.	required
output_model	type[BaseModel]	Pydantic model to validate against.	required

Returns:

Type	Description
BaseModel | None	BaseModel | None: Validated model instance, or None if no toolUse block was found or validation fails.

Source code in src/llmbo/adapters/nova.py

@classmethod
def validate_result(cls, result: dict[str, Any], output_model: type[BaseModel]) -> BaseModel | None:
    """Validate and parse output from Amazon Nova's Converse API.

    Searches the response content blocks for a ``toolUse`` entry
    and validates its input against the provided Pydantic model.

    Args:
        result (dict[str, Any]): Raw model output from the Converse
            API.
        output_model (type[BaseModel]): Pydantic model to validate
            against.

    Returns:
        BaseModel | None: Validated model instance, or None if no
            ``toolUse`` block was found or validation fails.
    """

    cls.logger.debug(f"Validating result against {output_model.__name__} schema")

    # Converse API returns data in: output -> message -> content -> toolUse
    output = result.get("output", {})
    message = output.get("message", {})
    content_blocks = message.get("content", [])

    for block in content_blocks:
        if "toolUse" in block:
            try:
                return output_model(**block["toolUse"]["input"])
            except ValidationError as e:
                cls.logger.debug(f"Validation failed: {e!s}")
                return None

    cls.logger.debug("Result contains no toolUse. The model hallucinated text.")
    return None

OpenAIAdapter

Bases: OpenAICompatibleAdapter

Adapter for OpenAI models in AWS Bedrock.

Source code in src/llmbo/adapters/openai_oss.py

class OpenAIAdapter(OpenAICompatibleAdapter):
    """Adapter for OpenAI models in AWS Bedrock."""

    logger = logging.getLogger(f"{__name__}.OpenAIAdapter")
    _provider_name = "OpenAI"

QwenAdapter

Bases: OpenAICompatibleAdapter

Adapter for Alibaba Qwen models in AWS Bedrock.

Source code in src/llmbo/adapters/qwen.py

class QwenAdapter(OpenAICompatibleAdapter):
    """Adapter for Alibaba Qwen models in AWS Bedrock."""

    logger = logging.getLogger(f"{__name__}.QwenAdapter")
    _provider_name = "Qwen"

StructuredBatchInferer

Bases: BatchInferer

A specialized BatchInferer that enforces structured outputs using Pydantic models.

Inspired by the instructor package, see: https://python.useinstructor.com/ This class extends BatchInferer to add schema validation and structured output handling using Pydantic models.

Parameters:

Name	Type	Description	Default
output_model	BaseModel	A Pydantic model defining the expected output structure	required
model_name	str	The name/ID of the AWS Bedrock model to use	required
bucket_name	str	The S3 bucket name for storing input/output data	required
region	str	The region to run the batch inference job in.	required
job_name	str	A unique name for the batch inference job	required
role_arn	str	The AWS IAM role ARN with necessary permissions	required
time_out_duration_hours	int	Maximum job runtime in hours. Defaults to 24.	24
session	Session	A boto3 session to be used for AWS API calls. If not provided, a new session will be created.	None

Source code in src/llmbo/structured_batch_inferer.py

class StructuredBatchInferer(BatchInferer):
    """A specialized BatchInferer that enforces structured outputs using Pydantic models.

    Inspired by the instructor package, see: https://python.useinstructor.com/
    This class extends BatchInferer to add schema validation and structured output
    handling using Pydantic models.

    Args:
        output_model (BaseModel): A Pydantic model defining the expected output structure
        model_name (str): The name/ID of the AWS Bedrock model to use
        bucket_name (str): The S3 bucket name for storing input/output data
        region (str): The region to run the batch inference job in.
        job_name (str): A unique name for the batch inference job
        role_arn (str): The AWS IAM role ARN with necessary permissions
        time_out_duration_hours (int, optional): Maximum job runtime in hours. Defaults to 24.
        session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                                           If not provided, a new session will be created.


    """

    logger = logging.getLogger(f"{__name__}.StructuredBatchInferer")

    def __init__(
        self,
        output_model: type[BaseModel],
        model_name: str,  # this should be an enum...
        bucket_name: str,
        region: str,
        job_name: str,
        role_arn: str,
        time_out_duration_hours: int = 24,
        session: boto3.Session | None = None,
        output_dir: str = ".",
    ):
        """Initialize a StructuredBatchInferer for schema-validated batch processing.

        Creates a batch inference manager that enforces structured outputs using
        a Pydantic model schema. Automatically configures the model to use tools
        for enforcing the output structure.

        Args:
            output_model (BaseModel): Pydantic model class defining the expected output structure
            model_name (str): The AWS Bedrock model identifier
            bucket_name (str): Name of the S3 bucket for storing job inputs and outputs
            region (str): Region of the LLM must match the bucket
            job_name (str): Unique identifier for this batch job
            role_arn (str): AWS IAM role ARN with permissions for Bedrock and S3 access
            time_out_duration_hours (int): Number of hours before the job times out
            session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                If not provided, a new session will be created.
            output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

        Raises:
            KeyError: If AWS_PROFILE environment variable is not set
            ValueError: If the provided role_arn doesn't exist or is invalid

        Example:
            >>> class PersonInfo(BaseModel):
            ...     name: str
            ...     age: int
            ...
            >>> sbi = StructuredBatchInferer(
            ...     output_model=PersonInfo,
            ...     model_name="anthropic.claude-3-haiku-20240307-v1:0",
            ...     bucket_name="my-inference-bucket",
            ...     job_name="structured-batch-2024",
            ...     role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole"
            ... )

        Note:
            - Converts the Pydantic model into a tool definition for the LLM
            - All results will be validated against the provided schema
            - Failed schema validations will raise errors during result processing
            - Inherits all base BatchInferer functionality
        """
        self.output_model = output_model

        self.logger.info(f"Initialized StructuredBatchInferer with {output_model.__name__} schema")

        super().__init__(
            model_name=model_name,
            bucket_name=bucket_name,
            region=region,
            job_name=job_name,
            role_arn=role_arn,
            time_out_duration_hours=time_out_duration_hours,
            session=session,
            output_dir=output_dir,
        )

    def prepare_requests(self, inputs: dict[str, ModelInput]):
        """Prepare structured batch inference requests with tool configurations.

        Extends the base preparation by adding tool definitions and tool choice
        parameters to each ModelInput. The tool definition is derived from the
        Pydantic output_model specified during initialization.

        Args:
            inputs (Dict[str, ModelInput]): Dictionary mapping record IDs to their corresponding
                ModelInput configurations. The record IDs will be used to track results.

        Raises:
            ValueError: If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

        Example:
            >>> class PersonInfo(BaseModel):
            ...     name: str
            ...     age: int
            >>> sbi = StructuredBatchInferer(output_model=PersonInfo, ...)
            >>> inputs = {
            ...     "001": ModelInput(
            ...         messages=[{"role": "user", "content": "John is 25 years old"}],
            ...     )
            ... }
            >>> sbi.prepare_requests(inputs)

        Note:
            - Automatically adds the output_model schema as a tool definition
            - Sets tool_choice to force use of the defined schema
            - Original ModelInputs are modified to include tool configurations
        """
        self.logger.info(f"Adding tool {self.output_model.__name__} to model input")
        self._check_input_length(inputs)
        for id, model_input in inputs.items():
            inputs[id] = self.adapter.prepare_model_input(model_input, self.output_model)

        self.requests = self._to_requests(inputs)

    def load_results(self):
        """Load and validate batch inference results against the output schema.

        Reads the output files downloaded from S3 and validates each result against
        the Pydantic output_model specified during initialization. Populates:
            - self.results: Raw inference results from the output JSONL file
            - self.manifest: Statistics about the job execution
            - self.instances: List of validated Pydantic model instances

        Raises:
            FileExistsError: If either the results or manifest files are not found locally
            ValueError: If any result fails schema validation or tool use validation

        Note:
            - Must call download_results() before calling this method
            - All results must conform to the specified output_model schema
            - Results must show successful tool use
        """
        super().load_results()
        self.instances = [
            {
                "recordId": result["recordId"],
                "outputModel": self.adapter.validate_result(result["modelOutput"], self.output_model),
            }
            if result.get("modelOutput")
            else None
            for result in self.results
        ]

    @classmethod
    def recover_details_from_job_arn(
        cls,
        job_arn: str,
        region: str,
        session: boto3.Session | None = None,
        output_dir: str = ".",
    ) -> "StructuredBatchInferer":
        """Placeholder method for interface consistency.

        This method exists to maintain compatibility with the parent class but
        is not implemented for structured jobs. Use `recover_structured_job`
        instead.

        Raises:
            NotImplementedError: Always raised when called.
        """
        raise NotImplementedError(
            "Cannot recover structured job without output_model. Use recover_structured_job instead."
        )

    @classmethod
    def recover_structured_job(
        cls,
        job_arn: str,
        region: str,
        output_model: type[BaseModel],
        session: boto3.Session | None = None,
        output_dir: str = ".",
    ) -> "StructuredBatchInferer":
        """Recover a StructuredBatchInferer instance from an existing job ARN.

        Used to reconstruct a StructuredBatchInferer object when the original Python
        process has terminated but the AWS job is still running or complete.

        Args:
            job_arn: (str) The AWS ARN of the existing batch inference job
            region: (str) the region where the job was scheduled
            output_model: (type[BaseModel]) A pydantic model describing the required output
            session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                                            If not provided, a new session will be created.
            output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

        Returns:
            StructuredBatchInferer: A configured instance with the job's details

        Raises:
            ValueError: If the job cannot be found or response is invalid

        Example:
            >>> job_arn = "arn:aws:bedrock:region:account:job/xyz123"
            >>> region = us-east-1"
            >>> sbi = StructuredBatchInferer.recover_structured_job(job_arn, region, some_model)
            >>> sbi.check_complete()
            'Completed'
        """
        cls.logger.info(f"Attempting to Recover BatchInferer from {job_arn}")
        session = session or boto3.Session()
        response = cls.check_for_existing_job(job_arn, region, session)

        try:
            # Extract required parameters from response
            job_name = response["jobName"]
            model_id = response["modelId"]
            bucket_name = response["inputDataConfig"]["s3InputDataConfig"]["s3Uri"].split("/")[2]
            role_arn = response["roleArn"]

            # Validate required files exist
            input_file = os.path.join(output_dir, f"{job_name}.jsonl")
            if not os.path.exists(input_file):
                cls.logger.error(f"Required input file not found: {input_file}")
                raise FileNotFoundError(f"Required input file not found: {input_file}")

            requests = cls._read_jsonl(input_file)

            sbi = cls(
                model_name=model_id,
                output_model=output_model,
                job_name=job_name,
                region=region,
                bucket_name=bucket_name,
                role_arn=role_arn,
                session=session,
                output_dir=output_dir,
            )
            sbi.job_arn = job_arn
            sbi.requests = requests
            sbi.job_status = response["status"]

            return sbi

        except (KeyError, IndexError) as e:
            cls.logger.error(f"Invalid job response format: {e!s}")
            raise ValueError(f"Invalid job response format: {e!s}") from e
        except Exception as e:
            cls.logger.error(f"Failed to recover job details: {e!s}")
            raise RuntimeError(f"Failed to recover job details: {e!s}") from e

__init__(output_model, model_name, bucket_name, region, job_name, role_arn, time_out_duration_hours=24, session=None, output_dir='.')

Initialize a StructuredBatchInferer for schema-validated batch processing.

Creates a batch inference manager that enforces structured outputs using a Pydantic model schema. Automatically configures the model to use tools for enforcing the output structure.

Parameters:

Name	Type	Description	Default
output_model	BaseModel	Pydantic model class defining the expected output structure	required
model_name	str	The AWS Bedrock model identifier	required
bucket_name	str	Name of the S3 bucket for storing job inputs and outputs	required
region	str	Region of the LLM must match the bucket	required
job_name	str	Unique identifier for this batch job	required
role_arn	str	AWS IAM role ARN with permissions for Bedrock and S3 access	required
time_out_duration_hours	int	Number of hours before the job times out	24
session	Session	A boto3 session to be used for AWS API calls. If not provided, a new session will be created.	None
output_dir	str	Directory for local JSONL files. Defaults to ".".	'.'

Raises:

Type	Description
KeyError	If AWS_PROFILE environment variable is not set
ValueError	If the provided role_arn doesn't exist or is invalid

Example

class PersonInfo(BaseModel): ... name: str ... age: int ... sbi = StructuredBatchInferer( ... output_model=PersonInfo, ... model_name="anthropic.claude-3-haiku-20240307-v1:0", ... bucket_name="my-inference-bucket", ... job_name="structured-batch-2024", ... role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole" ... )

Note

• Converts the Pydantic model into a tool definition for the LLM
• All results will be validated against the provided schema
• Failed schema validations will raise errors during result processing
• Inherits all base BatchInferer functionality

Source code in src/llmbo/structured_batch_inferer.py

def __init__(
    self,
    output_model: type[BaseModel],
    model_name: str,  # this should be an enum...
    bucket_name: str,
    region: str,
    job_name: str,
    role_arn: str,
    time_out_duration_hours: int = 24,
    session: boto3.Session | None = None,
    output_dir: str = ".",
):
    """Initialize a StructuredBatchInferer for schema-validated batch processing.

    Creates a batch inference manager that enforces structured outputs using
    a Pydantic model schema. Automatically configures the model to use tools
    for enforcing the output structure.

    Args:
        output_model (BaseModel): Pydantic model class defining the expected output structure
        model_name (str): The AWS Bedrock model identifier
        bucket_name (str): Name of the S3 bucket for storing job inputs and outputs
        region (str): Region of the LLM must match the bucket
        job_name (str): Unique identifier for this batch job
        role_arn (str): AWS IAM role ARN with permissions for Bedrock and S3 access
        time_out_duration_hours (int): Number of hours before the job times out
        session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
            If not provided, a new session will be created.
        output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

    Raises:
        KeyError: If AWS_PROFILE environment variable is not set
        ValueError: If the provided role_arn doesn't exist or is invalid

    Example:
        >>> class PersonInfo(BaseModel):
        ...     name: str
        ...     age: int
        ...
        >>> sbi = StructuredBatchInferer(
        ...     output_model=PersonInfo,
        ...     model_name="anthropic.claude-3-haiku-20240307-v1:0",
        ...     bucket_name="my-inference-bucket",
        ...     job_name="structured-batch-2024",
        ...     role_arn="arn:aws:iam::123456789012:role/BedrockBatchRole"
        ... )

    Note:
        - Converts the Pydantic model into a tool definition for the LLM
        - All results will be validated against the provided schema
        - Failed schema validations will raise errors during result processing
        - Inherits all base BatchInferer functionality
    """
    self.output_model = output_model

    self.logger.info(f"Initialized StructuredBatchInferer with {output_model.__name__} schema")

    super().__init__(
        model_name=model_name,
        bucket_name=bucket_name,
        region=region,
        job_name=job_name,
        role_arn=role_arn,
        time_out_duration_hours=time_out_duration_hours,
        session=session,
        output_dir=output_dir,
    )

load_results()

Load and validate batch inference results against the output schema.

Reads the output files downloaded from S3 and validates each result against the Pydantic output_model specified during initialization. Populates: - self.results: Raw inference results from the output JSONL file - self.manifest: Statistics about the job execution - self.instances: List of validated Pydantic model instances

Raises:

Type	Description
FileExistsError	If either the results or manifest files are not found locally
ValueError	If any result fails schema validation or tool use validation

Note

• Must call download_results() before calling this method
• All results must conform to the specified output_model schema
• Results must show successful tool use

Source code in src/llmbo/structured_batch_inferer.py

def load_results(self):
    """Load and validate batch inference results against the output schema.

    Reads the output files downloaded from S3 and validates each result against
    the Pydantic output_model specified during initialization. Populates:
        - self.results: Raw inference results from the output JSONL file
        - self.manifest: Statistics about the job execution
        - self.instances: List of validated Pydantic model instances

    Raises:
        FileExistsError: If either the results or manifest files are not found locally
        ValueError: If any result fails schema validation or tool use validation

    Note:
        - Must call download_results() before calling this method
        - All results must conform to the specified output_model schema
        - Results must show successful tool use
    """
    super().load_results()
    self.instances = [
        {
            "recordId": result["recordId"],
            "outputModel": self.adapter.validate_result(result["modelOutput"], self.output_model),
        }
        if result.get("modelOutput")
        else None
        for result in self.results
    ]

prepare_requests(inputs)

Prepare structured batch inference requests with tool configurations.

Extends the base preparation by adding tool definitions and tool choice parameters to each ModelInput. The tool definition is derived from the Pydantic output_model specified during initialization.

Parameters:

Name	Type	Description	Default
inputs	Dict[str, ModelInput]	Dictionary mapping record IDs to their corresponding ModelInput configurations. The record IDs will be used to track results.	required

Raises:

Type	Description
ValueError	If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

Example

class PersonInfo(BaseModel): ... name: str ... age: int sbi = StructuredBatchInferer(output_model=PersonInfo, ...) inputs = { ... "001": ModelInput( ... messages=[{"role": "user", "content": "John is 25 years old"}], ... ) ... } sbi.prepare_requests(inputs)

Note

• Automatically adds the output_model schema as a tool definition
• Sets tool_choice to force use of the defined schema
• Original ModelInputs are modified to include tool configurations

Source code in src/llmbo/structured_batch_inferer.py

def prepare_requests(self, inputs: dict[str, ModelInput]):
    """Prepare structured batch inference requests with tool configurations.

    Extends the base preparation by adding tool definitions and tool choice
    parameters to each ModelInput. The tool definition is derived from the
    Pydantic output_model specified during initialization.

    Args:
        inputs (Dict[str, ModelInput]): Dictionary mapping record IDs to their corresponding
            ModelInput configurations. The record IDs will be used to track results.

    Raises:
        ValueError: If len(inputs) < 100, as AWS Bedrock requires minimum batch size of 100

    Example:
        >>> class PersonInfo(BaseModel):
        ...     name: str
        ...     age: int
        >>> sbi = StructuredBatchInferer(output_model=PersonInfo, ...)
        >>> inputs = {
        ...     "001": ModelInput(
        ...         messages=[{"role": "user", "content": "John is 25 years old"}],
        ...     )
        ... }
        >>> sbi.prepare_requests(inputs)

    Note:
        - Automatically adds the output_model schema as a tool definition
        - Sets tool_choice to force use of the defined schema
        - Original ModelInputs are modified to include tool configurations
    """
    self.logger.info(f"Adding tool {self.output_model.__name__} to model input")
    self._check_input_length(inputs)
    for id, model_input in inputs.items():
        inputs[id] = self.adapter.prepare_model_input(model_input, self.output_model)

    self.requests = self._to_requests(inputs)

recover_details_from_job_arn(job_arn, region, session=None, output_dir='.') classmethod

Placeholder method for interface consistency.

This method exists to maintain compatibility with the parent class but is not implemented for structured jobs. Use recover_structured_job instead.

Raises:

Type	Description
NotImplementedError	Always raised when called.

Source code in src/llmbo/structured_batch_inferer.py

@classmethod
def recover_details_from_job_arn(
    cls,
    job_arn: str,
    region: str,
    session: boto3.Session | None = None,
    output_dir: str = ".",
) -> "StructuredBatchInferer":
    """Placeholder method for interface consistency.

    This method exists to maintain compatibility with the parent class but
    is not implemented for structured jobs. Use `recover_structured_job`
    instead.

    Raises:
        NotImplementedError: Always raised when called.
    """
    raise NotImplementedError(
        "Cannot recover structured job without output_model. Use recover_structured_job instead."
    )

recover_structured_job(job_arn, region, output_model, session=None, output_dir='.') classmethod

Recover a StructuredBatchInferer instance from an existing job ARN.

Used to reconstruct a StructuredBatchInferer object when the original Python process has terminated but the AWS job is still running or complete.

Parameters:

Name	Type	Description	Default
job_arn	str	(str) The AWS ARN of the existing batch inference job	required
region	str	(str) the region where the job was scheduled	required
output_model	type[BaseModel]	(type[BaseModel]) A pydantic model describing the required output	required
session	Session	A boto3 session to be used for AWS API calls. If not provided, a new session will be created.	None
output_dir	str	Directory for local JSONL files. Defaults to ".".	'.'

Returns:

Name	Type	Description
StructuredBatchInferer	StructuredBatchInferer	A configured instance with the job's details

Raises:

Type	Description
ValueError	If the job cannot be found or response is invalid

Example

job_arn = "arn:aws:bedrock:region:account:job/xyz123" region = us-east-1" sbi = StructuredBatchInferer.recover_structured_job(job_arn, region, some_model) sbi.check_complete() 'Completed'

Source code in src/llmbo/structured_batch_inferer.py

@classmethod
def recover_structured_job(
    cls,
    job_arn: str,
    region: str,
    output_model: type[BaseModel],
    session: boto3.Session | None = None,
    output_dir: str = ".",
) -> "StructuredBatchInferer":
    """Recover a StructuredBatchInferer instance from an existing job ARN.

    Used to reconstruct a StructuredBatchInferer object when the original Python
    process has terminated but the AWS job is still running or complete.

    Args:
        job_arn: (str) The AWS ARN of the existing batch inference job
        region: (str) the region where the job was scheduled
        output_model: (type[BaseModel]) A pydantic model describing the required output
        session (boto3.Session, optional): A boto3 session to be used for AWS API calls.
                                        If not provided, a new session will be created.
        output_dir (str, optional): Directory for local JSONL files. Defaults to ".".

    Returns:
        StructuredBatchInferer: A configured instance with the job's details

    Raises:
        ValueError: If the job cannot be found or response is invalid

    Example:
        >>> job_arn = "arn:aws:bedrock:region:account:job/xyz123"
        >>> region = us-east-1"
        >>> sbi = StructuredBatchInferer.recover_structured_job(job_arn, region, some_model)
        >>> sbi.check_complete()
        'Completed'
    """
    cls.logger.info(f"Attempting to Recover BatchInferer from {job_arn}")
    session = session or boto3.Session()
    response = cls.check_for_existing_job(job_arn, region, session)

    try:
        # Extract required parameters from response
        job_name = response["jobName"]
        model_id = response["modelId"]
        bucket_name = response["inputDataConfig"]["s3InputDataConfig"]["s3Uri"].split("/")[2]
        role_arn = response["roleArn"]

        # Validate required files exist
        input_file = os.path.join(output_dir, f"{job_name}.jsonl")
        if not os.path.exists(input_file):
            cls.logger.error(f"Required input file not found: {input_file}")
            raise FileNotFoundError(f"Required input file not found: {input_file}")

        requests = cls._read_jsonl(input_file)

        sbi = cls(
            model_name=model_id,
            output_model=output_model,
            job_name=job_name,
            region=region,
            bucket_name=bucket_name,
            role_arn=role_arn,
            session=session,
            output_dir=output_dir,
        )
        sbi.job_arn = job_arn
        sbi.requests = requests
        sbi.job_status = response["status"]

        return sbi

    except (KeyError, IndexError) as e:
        cls.logger.error(f"Invalid job response format: {e!s}")
        raise ValueError(f"Invalid job response format: {e!s}") from e
    except Exception as e:
        cls.logger.error(f"Failed to recover job details: {e!s}")
        raise RuntimeError(f"Failed to recover job details: {e!s}") from e

ToolChoice dataclass

Toolchoice details.

Source code in src/llmbo/models.py

@dataclass
class ToolChoice:
    """Toolchoice details."""

    type: Literal["any", "tool", "auto"]
    name: str | None = None