Source code for mmirage.core.process.batch.adapter

"""Provider-agnostic batch submission adapter contracts.

Adapters implement translation from internal request payloads into provider
request formats and normalize submission responses into a shared result shape.
"""

import abc
from dataclasses import dataclass, field
from typing import Any, Dict, Sequence, Tuple

from mmirage.config.batch_provider import BatchProviderConfig




@dataclass
class BatchSubmissionResult:
    """Normalized result returned by any provider adapter after chunk submission.

    Attributes:
        provider_batch_id: Provider-side identifier for the submitted job/batch.
        status: Provider status mapped to 'completed', 'failed', 'in_progress'
            or 'unknown', so generic code never reads a provider vocabulary.
        raw_response: Original provider response payload for traceability.
    """

    provider_batch_id: str
    status: str
    raw_response: Dict[str, Any] = field(default_factory=dict)





class BatchSubmissionAdapter(abc.ABC):
    """Abstract interface for provider-specific batch submission adapters.

    Implementations should be deterministic for request building and byte
    estimation so chunk boundaries can be reproduced across retries.
    """

    required_credentials: Tuple[str, ...] = tuple()



    @abc.abstractmethod
    def build_request(
        self,
        custom_id: str,
        payload: Dict[str, Any],
        config: BatchProviderConfig,
    ) -> Dict[str, Any]:
        """Build a single provider-ready request object.

        Args:
            custom_id: Stable request identifier used to map async results back
                to source rows.
            payload: Provider-neutral request payload assembled by the core
                processing layer.
            config: Provider configuration contract that may influence request
                shaping.

        Returns:
            A provider-specific request object represented as a mapping.
        """
        raise NotImplementedError()




    @abc.abstractmethod
    def estimate_request_bytes(self, request: Dict[str, Any]) -> int:
        """Estimate serialized UTF-8 bytes for a request payload.

        The estimate must match or safely upper-bound the size produced by the
        serializer used for submission so chunk boundaries are enforced
        correctly.

        Args:
            request: Provider request object returned by ``build_request``.

        Returns:
            Estimated byte size for the serialized request.
        """
        raise NotImplementedError()




    @abc.abstractmethod
    def submit_chunk(
        self,
        chunk_id: str,
        requests: Sequence[Dict[str, Any]],
        config: BatchProviderConfig,
    ) -> Dict[str, Any]:
        """Submit one pre-chunked request group to the provider.

        Args:
            chunk_id: Internal chunk identifier generated by orchestration.
            requests: Provider-ready request objects belonging to this chunk.
            config: Provider config containing credentials and submission knobs.

        Returns:
            Raw provider response payload as a mapping.
        """
        raise NotImplementedError()




    @abc.abstractmethod
    def parse_submission_result(
        self,
        raw_result: Dict[str, Any],
    ) -> BatchSubmissionResult:
        """Normalize provider submission output into a shared result model.

        Args:
            raw_result: Raw payload returned by ``submit_chunk``.

        Returns:
            A normalized ``BatchSubmissionResult`` for provider-neutral
            orchestration and metadata persistence.
        """
        raise NotImplementedError()




    @abc.abstractmethod
    def check_batch_status(
        self,
        provider_batch_id: str,
        config: BatchProviderConfig,
    ) -> BatchSubmissionResult:
        """Retrieve and normalize the latest status for a provider batch job.

        Args:
            provider_batch_id: Provider-side batch/job identifier to query.
            config: Provider configuration containing credentials and endpoint
                overrides.

        Returns:
            A normalized ``BatchSubmissionResult`` where ``status`` reflects the
            latest provider-reported lifecycle state for the batch.
        """
        raise NotImplementedError()




    @abc.abstractmethod
    def retrieve_results(
        self,
        provider_batch_id: str,
        config: BatchProviderConfig,
    ) -> Sequence[Dict[str, Any]]:
        """Download and parse completed batch results from the provider.

        Implementations should normalize each returned row into a plain mapping
        and, when a text payload is available, expose it as ``generated_text``
        so downstream collectors can consume a provider-agnostic result shape.
        Reported usage should likewise be exposed as ``input_tokens`` and
        ``output_tokens``, omitted when the provider reports none.

        Args:
            provider_batch_id: Provider-side batch/job identifier.
            config: Provider configuration containing credentials and endpoint
                overrides.

        Returns:
            Sequence of parsed result rows (provider JSONL records normalized to
            dictionaries) preserving provider output order.
        """
        raise NotImplementedError()