# main.py

from openai import OpenAI
import gradio as gr
from RAG_class import RAG_1177
from dotenv import load_dotenv
load_dotenv()

def predict(message, history):
    # Initialize RAG model
    rag = RAG_1177()
    
    # Prepare history in OpenAI format
    history_openai_format = []
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human })
        history_openai_format.append({"role": "assistant", "content": assistant})
    
    # System prompt
    system_prompt = """Du är en omtänksam assistent som ger sjukrådgivningsfrågor baserad dokumentation från 1177.se. Frågorna du kommer få kan delas i tre olika kategorier:
      1. Relevanta frågor där frågan kan svaras med dokumenten från 1177. När du svarar på dessa frågor ska du även ge källorna i slutet på svaret.
      2. Relevanta uppföljningsfrågor gällande ett svar. Svara med din bästa förmåga och du behöver INTE ge referenser.
      3. Relevanta frågor men som inte kan svaras med dokumenten från 1177. Då ska du tydligt säga att du inte kan svara på frågan utifrån 1177.se information och be användaren att omformulera sin fråga. DU ska INTE ge källor.
      4. Orelevanta frågor, då ber användaren att hålla sig till ämnet och du behöver inte förklara vad för information du har. Du ska INTE ge källor. """

    # Retrieve documents from RAG model
    documents, urls = rag.retrieve(message, 3)

    # User prompt
    user_prompt = f"svara frågan utifrån den givna informationen: \n fråga: {message} \n ----------------------------------- \n information från 1177.se:{documents} \n källor: {urls}"
    
    # Append prompts to history
    history_openai_format.append({"role": "system", "content": system_prompt})
    history_openai_format.append({"role": "user", "content": user_prompt})

    # Initialize OpenAI client
    client = OpenAI()

    # Generate response from OpenAI model
    response = client.chat.completions.create(
        model='gpt-3.5-turbo',
        messages=history_openai_format,
        temperature=0.5,
        stream=True
    )

    # Yield response progressively
    partial_message = ""
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            partial_message = partial_message + chunk.choices[0].delta.content
            yield partial_message

    yield partial_message

def main():
    examples = [
        "hur kan jag lindra min huvudvärk?",
        "Vad är symptomen för covid-19?",
        "Vilka läkemedel finns det för smärtlindring?",
        "Hur kan jag bäst lindra mina symtom på pollenallergi?",
        "Jag har svårt att somna på nätterna och känner mig trött under dagen. Vad kan jag göra för att förbättra min sömn?"
    ]

    description = "AI-chatbot som svarar på dina frågor utifrån dokument från 1177.se."

    gr.ChatInterface(predict, title="1177 AI-chatbot", examples=examples, description=description).launch(share=True)

if __name__ == "__main__":
    main()