from omegaconf import OmegaConf
from query import VectaraQuery
import os
import gradio as gr

def isTrue(x) -> bool:
    if isinstance(x, bool):
        return x
    return x.strip().lower() == 'true'

corpus_keys = str(os.environ['corpus_keys']).split(',')
cfg = OmegaConf.create({
    'corpus_keys': corpus_keys,
    'api_key': str(os.environ['api_key']),
    'title': os.environ['title'],
    'description': os.environ['description'],
    'source_data_desc': os.environ['source_data_desc'],
    'streaming': isTrue(os.environ.get('streaming', False)),
    'prompt_name': os.environ.get('prompt_name', None),
    'examples': os.environ.get('examples', None)
})

vq = VectaraQuery(cfg.api_key, cfg.corpus_keys, cfg.prompt_name)

def respond(message, history):
    if cfg.streaming:
        # Call stream response and stream output
        stream = vq.submit_query_streaming(message)
        
        
        outputs = ""
        for output in stream:
            outputs += output
            yield outputs
    else:
        # Call non-stream response and return message output
        response = vq.submit_query(message)
        yield response


cfg.description = f'''
                <table>
                  <tr>
                    <td style="width: 33%; vertical-align: bottom;"> <img src="https://github.com/david-oplatka/chatbot-streamlit/blob/main/Vectara-logo.png?raw=true"> </td>
                    <td style="width: 34%; vertical-align: middle;"> <h1>{cfg.title}</h1> </td>
                    <td style="width: 33%; vertical-align: bottom; text-align: left"> This demo uses Retrieval Augmented Generation from <a href="https://vectara.com/">Vectara</a><br>to ask questions about {cfg.source_data_desc}. </td>
                  </tr>
                </table>
                <center> <h2>{cfg.description}</h2></center>
                '''

css = """
table {
  border: none;
  width: 100%;
  table-layout: fixed;
  border-collapse: separate;
}
td {
  border: none;
  text-align: center;
}
img {
  width: 50%;
}
h1 {
  font-size: 3em; /* Adjust the size as needed */
}
"""

if cfg.examples:
    app_examples = [example.strip() for example in cfg.examples.split(",")]
else:
    app_examples = None

demo = gr.ChatInterface(respond, description = cfg.description, css = css,
                        chatbot = gr.Chatbot(value = [[None, "How may I help you?"]]), examples = app_examples)


if __name__ == "__main__":
    demo.launch()