#!/usr/bin/env python3
"""
Generated test file for Conversation ID LLM Test
Framework: litellm
Execution mode: sync
Streaming: no"""

import os
import sentry_sdk

import time

import litellm
from sentry_sdk.integrations.litellm import LiteLLMIntegration
from sentry_sdk.integrations.openai import OpenAIIntegration


sentry_sdk.init(
    dsn=os.environ.get("SENTRY_DSN"),
    traces_sample_rate=1.0,
    send_default_pii=True,
    stream_gen_ai_spans=True,
    integrations=[LiteLLMIntegration(include_prompts=True)],
    disabled_integrations=[OpenAIIntegration()],
)


def main():
    # Turn 1
    sentry_sdk.ai.set_conversation_id("conv-a")
    response = litellm.completion(
        model="anthropic/claude-haiku-4-5",
        messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant."
        },
        {
            "role": "user",
            "content": "What is the capital of France?"
        },
    ]
,
    )
    print(f"Turn 1 Response: {response.choices[0].message.content}")
    time.sleep(0.1)  # sleep is necessary for LiteLLM because it needs threaded callbacks to finish

    # Turn 2
    sentry_sdk.ai.set_conversation_id("conv-b")
    response = litellm.completion(
        model="anthropic/claude-haiku-4-5",
        messages=[
        {
            "role": "system",
            "content": "You are a math tutor."
        },
        {
            "role": "user",
            "content": "What is 2 + 2?"
        },
    ]
,
    )
    print(f"Turn 2 Response: {response.choices[0].message.content}")
    time.sleep(0.1)  # sleep is necessary for LiteLLM because it needs threaded callbacks to finish

    # Turn 3
    sentry_sdk.ai.set_conversation_id("conv-a")
    response = litellm.completion(
        model="anthropic/claude-haiku-4-5",
        messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant."
        },
        {
            "role": "user",
            "content": "What is the capital of France?"
        },
        {
            "role": "assistant",
            "content": "The capital of France is Paris."
        },
        {
            "role": "user",
            "content": "What about Germany?"
        },
    ]
,
    )
    print(f"Turn 3 Response: {response.choices[0].message.content}")
    time.sleep(0.1)  # sleep is necessary for LiteLLM because it needs threaded callbacks to finish

    # Turn 4
    sentry_sdk.ai.set_conversation_id("conv-b")
    response = litellm.completion(
        model="anthropic/claude-haiku-4-5",
        messages=[
        {
            "role": "system",
            "content": "You are a math tutor."
        },
        {
            "role": "user",
            "content": "What is 2 + 2?"
        },
        {
            "role": "assistant",
            "content": "2 + 2 equals 4."
        },
        {
            "role": "user",
            "content": "What about 3 + 3?"
        },
    ]
,
    )
    print(f"Turn 4 Response: {response.choices[0].message.content}")
    time.sleep(0.1)  # sleep is necessary for LiteLLM because it needs threaded callbacks to finish

if __name__ == "__main__":
    with sentry_sdk.start_transaction(op="test", name="Conversation ID LLM Test"):
        main()
    sentry_sdk.flush(timeout=5)