# Initialize and track the Bedrock client
bedrock_client = boto3.client("bedrock-runtime", region_name="us-east-1")
bedrock_client = track_bedrock(bedrock_client, project_name="bedrock-integration-demo")

MODEL_ID = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"

@track
def generate_story(prompt):
    res = bedrock_client.converse(
        modelId=MODEL_ID, 
        messages=[{"role": "user", "content": [{"text": prompt}]}],
        inferenceConfig={"temperature": 0.7, "maxTokens": 1000}
    )
    return res["output"]["message"]["content"][0]["text"]

@track
def generate_topic():
    prompt = "Generate a topic for a story about Opik."
    res = bedrock_client.converse(
        modelId=MODEL_ID, 
        messages=[{"role": "user", "content": [{"text": prompt}]}],
        inferenceConfig={"temperature": 0.7, "maxTokens": 500}
    )
    return res["output"]["message"]["content"][0]["text"]

@track
def generate_opik_story():
    topic = generate_topic()
    story = generate_story(topic)
    return story