from gradio_client import Client
from Classes.Owiki_Class import OWiki

class LocalLLM(OWiki):
    def __init__(self,**kwargs):
        self.gradio_url = kwargs['gradio_url']
        self.client = None
        try:
            self.client = Client(self.gradio_url)
        except:
            pass

    def predict(self, question, invocation_type, schemas):
        if self.client:
            response = self.client.predict(question, invocation_type, schemas)
            return response
        return None