def request_bedrock_data()

in Utils/llm/bedrock.py [0:0]


def request_bedrock_data(system_prompt, messages, model: Model):
    config = API[model]()

    system = [{"text": system_prompt}]
    formatted_messages = [
        {"role": message['role'], "content": [{"text": message['content']}]}
        for message in messages
    ]

    inf_params = {"temperature": temperature}

    response = client.converse(
        modelId=config["model_id"],
        messages=formatted_messages,
        system=system,
        inferenceConfig=inf_params,
    )

    return {
        'content': response["output"]["message"]["content"][0]["text"],
        'tokens': {
            "input_tokens": response["usage"]["inputTokens"],
            "output_tokens": response["usage"]["outputTokens"],
        }
    }