import gradio as gr
import requests
import os  # Import the os module to access environment variables

def chat_with_pdf(content):
    # Use the CHATPDFKEY environment variable for the API key
    api_key = os.getenv('CHATPDFKEY')
    if not api_key:
        return "API key (CHATPDFKEY) not set in environment variables."
    headers = {
        'x-api-key': api_key,
        "Content-Type": "application/json",
    }

    data = {
        'stream': True,
        'sourceId': 'src_REdjHAwIpm8Lhuk2HAtnm',  # Replace with your source ID
        'messages': [
            {
                'role': "user",
                'content': content,
            },
        ],
    }

    url = 'https://api.chatpdf.com/v1/chats/message'

    try:
        response = requests.post(url, json=data, headers=headers, stream=True)
        response.raise_for_status()

        if response.iter_content:
            max_chunk_size = 1023
            chat_response = ""
            for chunk in response.iter_content(max_chunk_size):
                text = chunk.decode()
                chat_response += text.strip()
            return chat_response
        else:
            return "No data received"
    except requests.exceptions.HTTPError as error:
        return f"HTTP Error: {error.response.status_code} - {error.response.text}"
    except requests.exceptions.RequestException as error:
        return f"Request Exception: {error}"

# Create Gradio interface
iface = gr.Interface(
    fn=chat_with_pdf,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Ask something..."),
    outputs="text",
    title="ChatPDF Query Interface",
    description="Type your question to get answers from PDFs. Make sure the CHATPDFKEY environment variable is set."
)

# Run the Gradio app
if __name__ == "__main__":
    iface.launch()