llmriddles/llms/llm_client.py

import time
import requests
import logging
import argparse


class LLMFlaskClient:
    def __init__(self, ip: str, port: int, max_retry: int = 3):
        self.ip = ip
        self.port = port

        self.url_prefix_format = 'http://{}:{}/'
        self.url = self.url_prefix_format.format(self.ip, self.port)
        self.max_retry = max_retry

        self.logger = logging.getLogger()
        self.logger.addHandler(logging.StreamHandler())
        self.logger.handlers[0].setFormatter(logging.Formatter("%(message)s"))

    def _request(self, name: str, data: dict):
        for _ in range(self.max_retry):
            try:
                self.logger.info(f'{name}\ndata: {data}')
                response = requests.post(self.url + name, json=data).json()
            except Exception as e:
                self.logger.warning('error: ', repr(e))
                time.sleep(1)
                continue
            if response['code'] == 0:
                return response['output']
            else:
                raise Exception(response['error_msg'])
        raise Exception("Web service failed. Please retry or contact with manager")

    def run(self, message: str) -> str:
        try:
            return self._request('ask_llm_for_answer', {'user_text': message})
        except Exception as e:
            return f"Error: {repr(e)}"


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', required=True)
    parser.add_argument('-p', '--port', required=True)
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()
    if args.debug:
        logging.getLogger().setLevel(logging.INFO)
    else:
        logging.getLogger().setLevel(logging.WARNING)

    client = LLMFlaskClient(args.ip, args.port)
    print(client.run('Please concatenate string "1+" and "1=3". Only give me the result without "".'))