app.py

import os
from typing import List
import argparse
from pokebot.rag import RAGApp, AssistantRole
from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import workflow, task

Traceloop.init(app_name="pokebot", disable_batch=False)

def _read_urls_from_file(filepath):
    # If seed is provided, parse and load the URLs in an array
    seed_list = []
    if filepath:
        with open(filepath, 'r') as file:
            seed_list = [line.strip() for line in file.readlines()]
    return seed_list

ASSISTANTS = {
    "healthcare": AssistantRole("Your Healthcare AI Assistant", 
                                _read_urls_from_file("data/med/articles/healthcare.txt"),
                                "diabetes"),
    "diabetes": AssistantRole("Your Diabetes AI Assistant", 
                                _read_urls_from_file("data/med/articles/diabetes.txt"),
                                "diabetes"),
    "default": AssistantRole("Your Diabetes Lite AI Assistant", 
                                ["https://www.niddk.nih.gov/health-information/diabetes/overview/what-is-diabetes"],
                                "diabetes"),
}


def _start_vuln_rag(args):
    assistant = ASSISTANTS.get(args.assistant_type, "default")
    # Initialize and run the RAGApp
    app = RAGApp(assistant=assistant)
    app.run()


def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description='RAG App with command-line options.')
    
    # Creating subparsers
    subparsers = parser.add_subparsers(title='command', dest='command', description='Choose a command to start the app.')
    
    # Subparser for the start command
    rag_parser = subparsers.add_parser('rag', help='Manage the RAG app')
    rag_subparser = rag_parser.add_subparsers(title='subcommand', dest='subcommand', description='Choose a sub command to start the app.')
    rag_start_parser = rag_subparser.add_parser("start", help="Start the rag app")
    rag_start_parser.add_argument('--assistant-type', type=str, 
                                  choices=["healthcare", "diabetes"], 
                                  default="default",
                                  help='Assistant Type')

    args = parser.parse_args()
    args.command = 'rag'
    args.subcommand = 'start'
    args.assistant_type = 'healthcare'

    if args.command == 'rag' and args.subcommand == 'start':
        _start_vuln_rag(args)
    else:
        print("Please specify a valid subcommand.")


if __name__ == "__main__":
   main()