TypeError: BaseLLM.predict() got an unexpected keyword argument 'context_str'

def query_google_llm(chat, query):
    response = chat.send_message(query)
    print(response.text)
    return response.text

chat = build_google_llm()

class PaLM(LLM):

    model_name = "Bard"
    total_tokens_used = 0
    last_token_usage = 0

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        print("prompt: ", prompt)
        response = query_google_llm(chat, prompt)

        print("response: ", response)
        return response

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"name_of_model": self.model_name}

    @property
    def _llm_type(self) -> str:
        return "custom"

QA_PROMPT_TMPL = (
    "### Human: Considering the following code:\n"
    "{context_str}\n"
    "{query_str}\n ### Assistant: \n"
)

def train(path):
    tokens = 0
    name = path.split("/")[-1]

    # get the documents inside the folder
    documents = SimpleDirectoryReader(path).load_data()
    print("Starting Vector construction at ", datetime.datetime.now())
    index = GPTSimpleVectorIndex.from_documents(documents)

    index.save_to_disk("indexes/" + name + ".json")

    return tokens

def query(query, toIndex):
    index = GPTSimpleVectorIndex.load_from_disk("indexes/" + toIndex + ".json")
    response = index.query(query)
    return response

response = query("question", "data")

INFO:llama_index.token_counter.token_counter:> [query] Total LLM token usage: 5002 tokens
INFO:llama_index.token_counter.token_counter:> [query] Total embedding token usage: 23 tokens

11:35
Local time: 30 mar 2023, 13:35
text-davinci, 2 requests
4,483 prompt + 512 completion = 4,995 tokens
11:35
Local time: 30 mar 2023, 13:35
text-embedding-ada-002-v2, 2 requests
56,906 prompt + 0 completion = 56,906 tokens