utils.py

import os
import json
from io import BytesIO

import boto3
import requests
from openai import OpenAI

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
BEDROCK_MODEL_ID = "anthropic.claude-3-5-sonnet-20240620-v1:0"
OPENAI_MODEL_ID = "o4-mini"
OPENAI_MODEL_IMAGE_MODEL_ID = "gpt-4.1-mini"

bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.getenv("AWS_REGION"))
openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


def invoke_bedrock_model(messages: list[dict]) -> str:
    """Invokes a Bedrock model with the provided messages."""
    response = bedrock_runtime.invoke_model(
        modelId=BEDROCK_MODEL_ID,
        body=json.dumps({"anthropic_version": "bedrock-2023-05-31", "max_tokens": 4096, "messages": messages}),
    )["body"].read()
    return json.loads(response)["content"][0]["text"]


def invoke_openai_model(messages: list[dict]) -> str:
    """Invokes an OpenAI model with the provided messages."""
    response = openai_client.responses.create(
        model=OPENAI_MODEL_IMAGE_MODEL_ID,
        input=messages,
    )
    return response.output_text


def get_file(task_id: str) -> BytesIO:
    """Fetches a file associated with a task ID from the default API URL.

    Parameters
    ----------
    task_id : str
        The ID of the task for which the file is to be fetched.

    Returns
    -------
    BytesIO
        A BytesIO object containing the file content.
    Raises
    ------
    requests.exceptions.RequestException
        If there is an error during the HTTP request.
    Exception
        For any other unexpected errors that may occur.
    """
    url = f"{DEFAULT_API_URL}/files/{task_id}"
    try:
        response = requests.get(url, timeout=15)
        response.raise_for_status()
        return BytesIO(response.content)
    except requests.exceptions.RequestException as e:
        print(f"Error fetching file for task {task_id}: {e}")
        raise
    except Exception as e:
        print(f"An unexpected error occurred fetching file for task {task_id}: {e}")
        raise


def s3_upload_file(file_content: BytesIO, bucket_name: str, object_name: str) -> None:
    """Uploads a file to an S3 bucket.

    Parameters
    ----------
    file_content : BytesIO
        The content of the file to upload.
    bucket_name : str
        The name of the S3 bucket.
    object_name : str
        The name of the object in the S3 bucket.

    Raises
    ------
    Exception
        If there is an error during the upload process.
    """
    try:
        s3_client = boto3.client("s3")
        s3_client.put_object(Bucket=bucket_name, Key=object_name, Body=file_content.getvalue())
    except Exception as e:
        print(f"Error uploading file to S3: {e}")
        raise


def s3_download_file(bucket_name: str, object_name: str) -> BytesIO:
    """Downloads a file from an S3 bucket.

    Parameters
    ----------
    bucket_name : str
        The name of the S3 bucket.
    object_name : str
        The name of the object in the S3 bucket.

    Returns
    -------
    BytesIO
        A BytesIO object containing the downloaded file content.

    Raises
    ------
    Exception
        If there is an error during the download process.
    """
    try:
        s3_client = boto3.client("s3")
        response = s3_client.get_object(Bucket=bucket_name, Key=object_name)
        return BytesIO(response["Body"].read())
    except Exception as e:
        print(f"Error downloading file from S3: {e}")
        raise