from langchain_community.document_loaders import PyPDFLoader
from langchain_core.messages import AIMessage, HumanMessage
from pydantic import BaseModel
import rag
import time
import gradio as gr
import requests
from main import run_server

class ChatInput(BaseModel):
    question: str
    
chat_history = []


def generate_response(chat_input: str, bot_message: str) -> str:
    url = "http://127.0.0.1:8000/generatechat/"
    payload = {
        'question': chat_input,
    }
    headers = {
        'Content-Type': 'application/json'
    }
    
    response = requests.post(url, json=payload, headers=headers)
    if response.status_code == 200:
        data = response.json()
        answer = data['response']['answer']
        print("Success:", response.json())
        
        # Get a typewriting animation response
        partial_response = ""
        for char in answer:
            partial_response += char
            yield partial_response
            time.sleep(0.005)
    else:
        print("Error:", response.status_code, response.text)
        return f"Error: {response.status_code}, {response.text}"
        
with gr.Blocks() as demo:
    with gr.Column():

        chatbot = gr.ChatInterface(
            fn=generate_response, 
            title="ThaiCodex Chat",
            description="Ask questions based on the content of the uploaded or specified PDF.",
        )

        # with gr.Row():
            # pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
            # upload_button = gr.Button("Load PDF")
        output_text = gr.Textbox(label="Status")
        # upload_button.click(, inputs=[pdf_input], outputs=output_text)

if __name__ == "__main__":
    demo.launch()
    run_server() # uvicorn api