#!/usr/bin/env python3
#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import sys

from gradio_client import Client
from rich.console import Console
from rich.markdown import Markdown

console = Console()
jarvis = Client("hadadrjt/ai")

model = "JARVIS: 2.1.2" # You can change the models here.

args = sys.argv[1:]
deep_search = False
if "-d" in args:
    deep_search = True
    args.remove("-d")
input_text = " ".join(args) if args else "Hi!"
jarvis.predict(new=model, api_name="/change_model")
if deep_search and model != "JARVIS: 2.1.2": # DO NOT CHANGE THIS!
    console.print("")
    console.print("------------------------------------------------------")
    console.print("[bold red]Deep Search is only available for model JARVIS: 2.1.2.[/bold red]")
    console.print("------------------------------------------------------")
    console.print("")
    deep_search = False
if deep_search:
    result = jarvis.predict(
        multi={"text": input_text},
        deep_search=True,
        api_name="/respond_async")
else:
    result = jarvis.predict(
        multi={"text": input_text},
        api_name="/api")
response_text = result[0][0][1]
console.print(Markdown(response_text))