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

from time import time
client = anthropic.Anthropic()

# Read file content into a string
with open("full_text.txt", "r") as file:
    text = file.read()

def claude_conversation(message, history):
    # 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": text,
                    "cache_control": {"type": "ephemeral"}
                }
            ],
            messages=messages
        )

        # Extract and return Claude's response

        print (response)

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

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

# Create the Gradio interface
demo = gr.ChatInterface(
    fn=claude_conversation,
    title="Claude Legal Document Analyzer",
    description="Ask questions about the legal document loaded from 'mainarbeit.txt'."
)

# Launch the app
demo.launch()