app.py

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

# import logging
# logging.basicConfig(level=logging.DEBUG)


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

corpus_ids = str(os.environ['corpus_ids']).split(',')
cfg = OmegaConf.create({
    'customer_id': str(os.environ['customer_id']),
    'corpus_ids': corpus_ids,
    '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)
})


vq = VectaraQuery(cfg.api_key, cfg.customer_id, cfg.corpus_ids, cfg.prompt_name)


def respond(message, history):
    if cfg.streaming:
        # Call stream response and stream output
        # logging.debug(f'Streaming query called with message {message}')
        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.title = f'''<center> <img src="https://github.com/david-oplatka/chatbot-streamlit/blob/main/Vectara-logo.png?raw=true" width="200px" height="40px">
                <h1>{cfg.title}</h1> </center>
                '''

cfg.description = f'''<center> <h2>{cfg.description}</h2>
                      <br>
                      This demo uses Retrieval Augmented Generation to ask questions about {cfg.source_data_desc}</center>
                      '''

demo = gr.ChatInterface(respond, title = cfg.title, description = cfg.description,
                        chatbot = gr.Chatbot(value = [[None, "How may I help you?"]], scale=3), examples = ["Hello", "Hi", "How are you?"])


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