import gradio as gr
import anthropic
from anthropic.types import ContentBlock
from time import time

def claude_conversation(message, history, api_key, uploaded_file):
    if not api_key:
        return "Please enter your Anthropic API key."

    if uploaded_file is None:
        return "Please upload a text document."

    client = anthropic.Anthropic(api_key=api_key)
    
    # Read the uploaded file content
    try:
        if hasattr(uploaded_file, 'read'):
            text = uploaded_file.read().decode('utf-8')
        else:
            # If it's already a string (filename), we need to open and read the file
            with open(uploaded_file, 'r') as file:
                text = file.read()
    except Exception as e:
        return f"Error reading file: {str(e)}"
    
    # Prepare the conversation history
    messages = []
    for human, assistant in history:
        if human.strip():  # Only add non-empty human messages
            messages.append({"role": "user", "content": human})
        if assistant.strip():  # Only add non-empty assistant messages
            messages.append({"role": "assistant", "content": assistant})
    
    # Add the new message if it's not empty
    if message.strip():
        messages.append({"role": "user", "content": message})
    else:
        return "Please enter a non-empty message."

    try:
        # Make the API call
        start = time()
        response = client.beta.prompt_caching.messages.create(
            model="claude-3-5-sonnet-20240620",
            max_tokens=1024,
            system=[
                {
                    "type": "text",
                    "text": "You are an AI assistant tasked with analyzing legal documents."
                },
                {
                    "type": "text",
                    "text": f"Document content:\n\n{text}",
                    "cache_control": {"type": "ephemeral"}
                }
            ],
            messages=messages
        )

        end = time()
        print(f"Elapsed time: {end - start} seconds")

        print (response)
        return response.content[0].text
    except anthropic.APIError as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Claude Prompt Caching Demo")
    gr.Markdown("Upload a lengthy document and ask questions about it.")
    
    api_key_input = gr.Textbox(label="Enter your Anthropic API key", type="password")
    file_upload = gr.File(label="Upload text document", file_types=[".txt"])
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    submit_button = gr.Button("Submit")
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history, api_key, uploaded_file):
        bot_message = claude_conversation(message, chat_history, api_key, uploaded_file)
        chat_history.append((message, bot_message))
        return "", chat_history

    submit_button.click(respond, inputs=[msg, chatbot, api_key_input, file_upload], outputs=[msg, chatbot])
    msg.submit(respond, inputs=[msg, chatbot, api_key_input, file_upload], outputs=[msg, chatbot])

# Launch the app
demo.launch()