Function retrieve_batch_results

scrapegraphai/utils/batch_api.py:250–316

Retrieve and parse results from a completed batch. Args: client: An initialized OpenAI client. batch_info: A BatchJobInfo from a completed batch. Returns: List of BatchResult objects, one per request, ordered by their custom_id.

(
    client: OpenAI,
    batch_info: BatchJobInfo,
)

def retrieve_batch_results(
    client: OpenAI,
    batch_info: BatchJobInfo,
) -> List[BatchResult]:
    """Retrieve and parse results from a completed batch.

    Args:
        client: An initialized OpenAI client.
        batch_info: A BatchJobInfo from a completed batch.

    Returns:
        List of BatchResult objects, one per request,
        ordered by their custom_id.
    """
    if not batch_info.output_file_id:
        raise ValueError(
            f"Batch {batch_info.batch_id} has no output file. "
            f"Status: {batch_info.status}"
        )

    logger.info(f"Downloading results from {batch_info.output_file_id}...")

    output_content = client.files.content(batch_info.output_file_id).text
    results = []

    for line in output_content.strip().split("\n"):
        if not line:
            continue

        response_data = json.loads(line)
        custom_id = response_data["custom_id"]

        error = response_data.get("error")
        if error:
            results.append(BatchResult(
                custom_id=custom_id,
                error=json.dumps(error),
            ))
            continue

        body = response_data.get("response", {}).get("body", {})
        choices = body.get("choices", [])

        if choices:
            content = choices[0].get("message", {}).get("content", "")
            usage = body.get("usage")
            results.append(BatchResult(
                custom_id=custom_id,
                content=content,
                usage=usage,
            ))
        else:
            results.append(BatchResult(
                custom_id=custom_id,
                error="No choices returned in response",
            ))

    # Sort by custom_id to maintain order