Spaces:
Running
Running
File size: 2,519 Bytes
3095540 724afe3 3095540 724afe3 3095540 724afe3 7ec877b 724afe3 7ec877b 724afe3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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()
|