from sawsi.aws import shared
import json


class Bedrock:
    def __init__(self, boto3_session, region=shared.DEFAULT_REGION):
        self.client = boto3_session.client('bedrock-runtime', region_name=region)
        self.region = region

    def invoke_model(self, model_id, body, content_type='application/json', accept='application/json'):
        response = self.client.invoke_model(
            modelId=model_id,
            body=json.dumps(body) if isinstance(body, dict) else body,
            contentType=content_type,
            accept=accept,
        )
        return json.loads(response['body'].read())

    def converse(self, model_id, messages, system=None, inference_config=None, tool_config=None):
        params = {
            'modelId': model_id,
            'messages': messages,
        }
        if system:
            params['system'] = system
        if inference_config:
            params['inferenceConfig'] = inference_config
        if tool_config:
            params['toolConfig'] = tool_config
        response = self.client.converse(**params)
        return response

    def converse_stream(self, model_id, messages, system=None, inference_config=None, tool_config=None):
        params = {
            'modelId': model_id,
            'messages': messages,
        }
        if system:
            params['system'] = system
        if inference_config:
            params['inferenceConfig'] = inference_config
        if tool_config:
            params['toolConfig'] = tool_config
        response = self.client.converse_stream(**params)
        return response
