import gradio as gr
import ollama
# import asyncio
# from openai import AsyncOpenAI
from RAGModule import RAGModule

# Instantiate the RAG module
ollama.pull("MarcoAland/llama3.1-rag.indo")
RAG_Trwira = RAGModule()

# # Configure the async OpenAI client
# client = AsyncOpenAI(api_key="34.69.9.203", base_url="http://34.69.9.203:11434/v1")

# settings = {
#     "model": "MarcoAland/llama3.1-rag-indo",
#     "temperature": 0.3,
#     "max_tokens": 2048,
# }

# async def generate_response(user_input: str) -> str:
#     message = "Namamu adalah Mitrakara.\n\n" + user_input
#     # Call documents options or not
#     if "dokumen" in message.lower() or "document" in message.lower() or "documents" in message.lower():
#         prompt = RAG_Trwira.main(message[10:])
#     else:
#         prompt = message
        
#     # Format the messages as a list of message dictionaries
#     message_formated = [
#         {"role": "user", "content": prompt}
#     ]

#     # Use streaming to handle partial responses
#     stream = await client.chat.completions.create(messages=message_formated, stream=True, **settings)
    
#     response = ""
#     async for part in stream:
#         if token := part.choices[0].delta.content or "":
#             response += token

#     return response

# def chat(user_input: str):
#     # Call the asynchronous response generation function
#     response = asyncio.run(generate_response(user_input))
#     return response

def chat(message: str, chat_history: str):
    if "dokumen" in message.lower() or "document" in message.lower() or "documents" in message.lower():
        prompt = RAG_Trwira.main(message[10:])
    else:
        prompt = message

    stream = ollama.chat(
                            model='MarcoAland/llama3.1-rag-indo',
                            messages=[{'role': 'user', 'content': prompt}],
                            stream=True,
                        )

    response_text = ''
    for chunk in stream:
        response_text += chunk['message']['content']
        yield response_text

# Define the Gradio interface
iface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(label="Masukkan pertanyaan anda", placeholder="Tanyakan saja padaku🌟"),
    outputs=gr.Textbox(label="Respons Mitrakara"),
    title="Hai, namaku Mitrakara. Selamat datang!👋",
    description="Berikut adalah beberapa tips untuk bertanya denganku✨✨✨\n1. Gunakan kata 'document:' jika ingin bertanya mengenai dokumen/administrasi perusahaan.\n2. Gunakan kalimat tanya yang baik.\n3. Enjoy the conversation.😊"
)

# Launch the Gradio interface
if __name__ == "__main__":
    iface.launch(share=False)