diff --git a/src/[44 b/src/[44 new file mode 100644 index 0000000..e69de29 diff --git a/src/aios/storage/storage.py b/src/aios/storage/storage.py index fe0b3fd..11edd29 100644 --- a/src/aios/storage/storage.py +++ b/src/aios/storage/storage.py @@ -216,6 +216,13 @@ class AIStorage: ~/myai/ """ return Path.home() / "myai" + + def get_download_dir(self) -> str: + """ + download dir is the dir for user to store the files downloaded with the system. + ~/myai/download + """ + return f"{self.get_myai_dir()}/download" def get_db(self,app_name:str)->ResourceLocation: pass @@ -242,5 +249,3 @@ class AIStorage: except Exception as e: logger.error(f"open or create file {path} failed! {str(e)}") - - diff --git a/src/aios_kernel/code_interpreter.py b/src/aios_kernel/code_interpreter.py new file mode 100644 index 0000000..850356a --- /dev/null +++ b/src/aios_kernel/code_interpreter.py @@ -0,0 +1,426 @@ +import logging +import os +import pathlib +import shutil +import subprocess +import sys +import re +import time +import ast +from concurrent.futures import ThreadPoolExecutor +from hashlib import md5 +from typing import Optional, Union, List, Tuple +from generic_escape import GenericEscape + +from aios_kernel import AIStorage + +try: + import docker +except ImportError: + docker = None + +CODE_BLOCK_PATTERN = r"```[ \t]*(\w+)?[ \t]*\r?\n(.*?)\r?\n[ \t]*```" +UNKNOWN = "unknown" +TIMEOUT_MSG = "Timeout" +DEFAULT_TIMEOUT = 600 +WIN32 = sys.platform == "win32" +PATH_SEPARATOR = WIN32 and "\\" or "/" + +logger = logging.getLogger(__name__) + + +BUILT_IN_MODULES = set( + [ + "sys", + "os", + "math", + "random", + "datetime", + "json", + "re", + "subprocess", + "time", + "threading", + "logging", + "collections", + "itertools", + "functools", + "operator", + "pathlib", + "shutil", + "tempfile", + "pickle", + "io", + "argparse", + "typing", + "unittest", + "contextlib", + "abc", + "heapq", + "bisect", + "copy", + "decimal", + "fractions", + "hashlib", + "secrets", + "statistics", + "difflib", + "doctest", + "enum", + "inspect", + "traceback", + "weakref", + "gc", + "mmap", + "msvcrt", + "winreg", + "array", + "audioop", + "binascii", + "cProfile", + "concurrent.futures", + "configparser", + "csv", + "ctypes", + "dateutil", + "dis", + "fnmatch", + "getopt", + "glob", + "gzip", + "pdb", + "pprint", + "profile", + "pstats", + "queue", + "socket", + "sqlite3", + "ssl", + "struct", + "tarfile", + "telnetlib", + "timeit", + "tokenize", + "uuid", + "xml", + "zipfile", + "zlib", + ] +) + + +def get_imports(code: str) -> List[str]: + root = ast.parse(code) + + imports = [] + for node in ast.iter_child_nodes(root): + if isinstance(node, ast.Import): + module_names = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + module_names = [node.module] + else: + continue + + for name in module_names: + # Exclude built-in modules + if name not in BUILT_IN_MODULES: + imports.append(name) + + return imports + + +def write_requirements(code: str, requirements_filepath: str): + imports = get_imports(code) + + with open(requirements_filepath, "w") as file: + for module in imports: + file.write(module + "\n") + + +def _cmd(lang): + if lang.startswith("python") or lang in ["bash", "sh", "powershell"]: + return lang + if lang in ["shell"]: + return "sh" + if lang in ["ps1"]: + return "powershell" + raise NotImplementedError(f"{lang} not recognized in code execution") + + +def create_runner(code: str, timeout: int = 30) -> str: + """ + Create a Python script that runs the code and prints the output + """ + code = GenericEscape().escape(code) + # Create a runner script + runner = f""" +import os +import subprocess + +my_env = os.environ.copy() +my_env["PYTHONIOENCODING"] = "utf-8" + +process = subprocess.Popen( + f"python -i -q -u".split(), + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=0, + universal_newlines=True, + env=my_env +) + +process.stdin.write("{code}" + "\\n") +process.stdin.write("exit()\\n") +process.stdin.flush() + +try: + process.wait({timeout}) +except Exception as e: + process.terminate() + +for line in iter(process.stdout.readline, ""): + print(line) + +for line in iter(process.stderr.readline, ""): + if line.startswith(">>>"): + continue + print(line) +""" + return runner + + +def _run_cmd(cmd: [str], work_dir: str, timeout: int) -> str: + if WIN32: + logger.warning("SIGALRM is not supported on Windows. No timeout will be enforced.") + result = subprocess.run( + cmd, + cwd=work_dir, + capture_output=True, + text=True, + ) + else: + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + subprocess.run, + cmd, + cwd=work_dir, + capture_output=True, + text=True, + ) + result = future.result(timeout=timeout) + return result + + +def execute_code( + code: Optional[str] = None, + timeout: Optional[int] = None, + filename: Optional[str] = None, + work_dir: Optional[str] = None, + use_docker: Optional[Union[List[str], str, bool]] = None, + lang: Optional[str] = "python", +) -> Tuple[int, str]: + """Execute code in a docker container. + This function is not tested on MacOS. + + Args: + code (Optional, str): The code to execute. + If None, the code from the file specified by filename will be executed. + Either code or filename must be provided. + timeout (Optional, int): The maximum execution time in seconds. + If None, a default timeout will be used. The default timeout is 600 seconds. On Windows, the timeout is not enforced when use_docker=False. + filename (Optional, str): The file name to save the code or where the code is stored when `code` is None. + If None, a file with a randomly generated name will be created. + The randomly generated file will be deleted after execution. + The file name must be a relative path. Relative paths are relative to the working directory. + work_dir (Optional, str): The working directory for the code execution. + If None, a default working directory will be used. + The default working directory is the "extensions" directory under + "path_to_autogen". + use_docker (Optional, list, str or bool): The docker image to use for code execution. + If a list or a str of image name(s) is provided, the code will be executed in a docker container + with the first image successfully pulled. + If None, False or empty, the code will be executed in the current environment. + Default is None, which will be converted into an empty list when docker package is available. + Expected behaviour: + - If `use_docker` is explicitly set to True and the docker package is available, the code will run in a Docker container. + - If `use_docker` is explicitly set to True but the Docker package is missing, an error will be raised. + - If `use_docker` is not set (i.e., left default to None) and the Docker package is not available, a warning will be displayed, but the code will run natively. + If the code is executed in the current environment, + the code must be trusted. + lang (Optional, str): The language of the code. Default is "python". + + Returns: + int: 0 if the code executes successfully. + str: The error message if the code fails to execute; the stdout otherwise. + """ + if all((code is None, filename is None)): + error_msg = f"Either {code=} or {filename=} must be provided." + logger.error(error_msg) + raise AssertionError(error_msg) + + # Warn if use_docker was unspecified (or None), and cannot be provided (the default). + # In this case the current behavior is to fall back to run natively, but this behavior + # is subject to change. + if use_docker is None: + if docker is None: + use_docker = False + logger.warning( + "execute_code was called without specifying a value for use_docker. Since the python docker package is not available, code will be run natively. Note: this fallback behavior is subject to change" + ) + else: + # Default to true + use_docker = True + + timeout = timeout or DEFAULT_TIMEOUT + original_filename = filename + if WIN32 and lang in ["sh", "shell"] and (not use_docker): + lang = "ps1" + if filename is None: + code_hash = md5(code.encode()).hexdigest() + # create a file with a automatically generated name + filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}" + if work_dir is None: + WORKING_DIR = os.path.join(AIStorage.get_instance().get_myai_dir(), "tmp_code") + pathlib.Path(WORKING_DIR).mkdir(exist_ok=True) + work_dir = os.path.join(WORKING_DIR, code_hash) + pathlib.Path(work_dir).mkdir(exist_ok=True) + filepath = os.path.join(work_dir, filename) + file_dir = os.path.dirname(filepath) + os.makedirs(file_dir, exist_ok=True) + if code is not None: + write_requirements(code, os.path.join(file_dir, "requirements.txt")) + code = create_runner(code, 30) + with open(filepath, "w", encoding="utf-8") as fout: + fout.write(code) + + + # check if already running in a docker container + in_docker_container = os.path.exists("/.dockerenv") + if not use_docker or in_docker_container: + try: + env_cmd = ["python", "-m", "venv", os.path.join(file_dir, "venv")] + _run_cmd(env_cmd, file_dir, timeout) + if WIN32: + venv_path = os.path.join(file_dir, "venv", "Scripts") + else: + venv_path = os.path.join(file_dir, "venv", "bin") + pip_cmd = [os.path.join(venv_path, "python"), "-m", "pip", "install", "-r", "requirements.txt"] + _run_cmd(pip_cmd, file_dir, timeout) + # already running in a docker container + cmd = [ + os.path.join(venv_path, "python"), + f".\\{filename}" if WIN32 else filename, + ] + result = _run_cmd(cmd, file_dir, timeout) + except TimeoutError: + if original_filename is None: + shutil.rmtree(os.path.join(file_dir, "venv")) + os.remove(filepath) + os.remove(os.path.join(file_dir, "requirements.txt")) + try: + os.removedirs(file_dir) + except Exception: + pass + return 1, TIMEOUT_MSG + if original_filename is None: + shutil.rmtree(os.path.join(file_dir, "venv")) + os.remove(filepath) + os.remove(os.path.join(file_dir, "requirements.txt")) + try: + os.removedirs(file_dir) + except Exception: + pass + if result.returncode: + logs = result.stderr + if original_filename is None: + abs_path = str(pathlib.Path(filepath).absolute()) + logs = logs.replace(str(abs_path), "").replace(filename, "") + else: + abs_path = str(pathlib.Path(work_dir).absolute()) + PATH_SEPARATOR + logs = logs.replace(str(abs_path), "") + else: + logs = result.stdout + return result.returncode, logs + + # create a docker client + client = docker.from_env() + image_list = ( + ["python:3-alpine", "python:3", "python:3-windowsservercore"] + if use_docker is True + else [use_docker] + if isinstance(use_docker, str) + else use_docker + ) + for image in image_list: + # check if the image exists + try: + client.images.get(image) + break + except docker.errors.ImageNotFound: + # pull the image + logger.info("Pulling image", image) + try: + client.images.pull(image, stream=True, decode=True) + break + except docker.errors.DockerException as e: + logger.error("Failed to pull image", image) + logger.exception(e) + # get a randomized str based on current time to wrap the exit code + exit_code_str = f"exitcode{time.time()}" + start_str = f'start{time.time()}' + abs_path = pathlib.Path(work_dir).absolute() + cmd = [ + "sh", + "-c", + f"pip install --quiet -r requirements.txt; echo -n {start_str}; {_cmd(lang)} {filename}; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str};", + ] + # create a docker container + container = client.containers.run( + image, + command=cmd, + working_dir="/workspace", + detach=True, + # get absolute path to the working directory + volumes={abs_path: {"bind": "/workspace", "mode": "rw"}}, + ) + start_time = time.time() + while container.status != "exited" and time.time() - start_time < timeout: + # Reload the container object + container.reload() + if container.status != "exited": + container.stop() + container.remove() + if original_filename is None: + os.remove(filepath) + return 1, TIMEOUT_MSG, image + # get the container logs + logs: str = container.logs().decode("utf-8").rstrip() + start_pos = logs.find(start_str) + if start_pos != -1: + logs = logs[start_pos + len(start_str):] + # # commit the image + # tag = filename.replace("/", "") + # container.commit(repository="python", tag=tag) + # remove the container + container.remove() + # check if the code executed successfully + exit_code = container.attrs["State"]["ExitCode"] + if exit_code == 0: + # extract the exit code from the logs + pattern = re.compile(f"{exit_code_str}(\\d+){exit_code_str}") + match = pattern.search(logs) + exit_code = 1 if match is None else int(match.group(1)) + # remove the exit code from the logs + logs = logs if match is None else pattern.sub("", logs) + + if original_filename is None: + os.remove(filepath) + os.remove(os.path.join(file_dir, "requirements.txt")) + os.removedirs(file_dir) + if exit_code: + logs = logs.replace(f"/workspace/{filename if original_filename is None else ''}", "") + # return the exit code, logs and image + return exit_code, logs + diff --git a/src/aios_kernel/code_interpreter_function.py b/src/aios_kernel/code_interpreter_function.py new file mode 100644 index 0000000..0bf366c --- /dev/null +++ b/src/aios_kernel/code_interpreter_function.py @@ -0,0 +1,41 @@ +from typing import Dict + +from aios_kernel.ai_function import AIFunction +from aios_kernel.code_interpreter import execute_code + + +class CodeInterpreterFunction(AIFunction): + def __init__(self): + self.func_id = "code_interpreter" + self.description = "execute python code" + + def get_name(self) -> str: + return self.func_id + + def get_description(self) -> str: + return self.description + + def get_parameters(self) -> Dict: + return { + "type": "object", + "properties": { + "code": {"type": "string", "description": "python code"} + } + } + + async def execute(self, **kwargs) -> str: + code = kwargs.get("code") + ret_code, result = execute_code(code=code) + if ret_code == 0: + return result.strip() + else: + return result.strip() + + def is_local(self) -> bool: + return True + + def is_in_zone(self) -> bool: + return True + + def is_ready_only(self) -> bool: + return False diff --git a/src/aios_kernel/duckduckgo_text_search_function.py b/src/aios_kernel/duckduckgo_text_search_function.py new file mode 100644 index 0000000..154f761 --- /dev/null +++ b/src/aios_kernel/duckduckgo_text_search_function.py @@ -0,0 +1,52 @@ +import json +from typing import Dict + +from aios_kernel.ai_function import AIFunction +from duckduckgo_search import AsyncDDGS + + +class DuckDuckGoTextSearchFunction(AIFunction): + def __init__(self): + self.name = "duckduckgo_text_search" + self.description = "Search text from duckduckgo.com" + self.region = "wt-wt" + self.safesearch = "moderate" + self.time = "y" + self.max_results = 5 + + def get_name(self) -> str: + return self.name + + def get_description(self) -> str: + return self.description + + def get_parameters(self) -> Dict: + return {"type": "object", + "properties": { + "query": {"type": "string", "description": "The query to search for."} + } + } + + async def execute(self, **kwargs) -> str: + query = kwargs.get("query") + + async with AsyncDDGS() as ddgs: + results = [r async for r in ddgs.text( + query, + region=self.region, + safesearch=self.safesearch, + timelimit=self.time, + backend="api", + max_results=self.max_results + )] + + return json.dumps(results) + + def is_local(self) -> bool: + return True + + def is_in_zone(self) -> bool: + return True + + def is_ready_only(self) -> bool: + return False diff --git a/src/aios_kernel/sql_database.py b/src/aios_kernel/sql_database.py new file mode 100644 index 0000000..efee9b4 --- /dev/null +++ b/src/aios_kernel/sql_database.py @@ -0,0 +1,493 @@ +""" +Taken from: langchain +SQLAlchemy wrapper around a database. +""" +from __future__ import annotations +import os + +import warnings +from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Union + +import sqlalchemy +from sqlalchemy import MetaData, Table, create_engine, inspect, select, text +from sqlalchemy.engine import Engine +from sqlalchemy.exc import ProgrammingError, SQLAlchemyError +from sqlalchemy.schema import CreateTable +from sqlalchemy.types import NullType + + +def get_from_env(key: str, env_key: str, default: Optional[str] = None) -> str: + """Get a value from a dictionary or an environment variable.""" + if env_key in os.environ and os.environ[env_key]: + return os.environ[env_key] + elif default is not None: + return default + else: + raise ValueError( + f"Did not find {key}, please add an environment variable" + f" `{env_key}` which contains it, or pass" + f" `{key}` as a named parameter." + ) + + +def _format_index(index: sqlalchemy.engine.interfaces.ReflectedIndex) -> str: + return ( + f'Name: {index["name"]}, Unique: {index["unique"]},' + f' Columns: {str(index["column_names"])}' + ) + + +def truncate_word(content: Any, *, length: int, suffix: str = "...") -> str: + """ + Truncate a string to a certain number of words, based on the max string + length. + """ + + if not isinstance(content, str) or length <= 0: + return content + + if len(content) <= length: + return content + + return content[: length - len(suffix)].rsplit(" ", 1)[0] + suffix + + +class SQLDatabase: + """SQLAlchemy wrapper around a database.""" + + def __init__( + self, + engine: Engine, + schema: Optional[str] = None, + metadata: Optional[MetaData] = None, + ignore_tables: Optional[List[str]] = None, + include_tables: Optional[List[str]] = None, + sample_rows_in_table_info: int = 3, + indexes_in_table_info: bool = False, + custom_table_info: Optional[dict] = None, + view_support: bool = False, + max_string_length: int = 300, + ): + """Create engine from database URI.""" + self._engine = engine + self._schema = schema + if include_tables and ignore_tables: + raise ValueError("Cannot specify both include_tables and ignore_tables") + + self._inspector = inspect(self._engine) + + # including view support by adding the views as well as tables to the all + # tables list if view_support is True + self._all_tables = set( + self._inspector.get_table_names(schema=schema) + + (self._inspector.get_view_names(schema=schema) if view_support else []) + ) + + self._include_tables = set(include_tables) if include_tables else set() + if self._include_tables: + missing_tables = self._include_tables - self._all_tables + if missing_tables: + raise ValueError( + f"include_tables {missing_tables} not found in database" + ) + self._ignore_tables = set(ignore_tables) if ignore_tables else set() + if self._ignore_tables: + missing_tables = self._ignore_tables - self._all_tables + if missing_tables: + raise ValueError( + f"ignore_tables {missing_tables} not found in database" + ) + usable_tables = self.get_usable_table_names() + self._usable_tables = set(usable_tables) if usable_tables else self._all_tables + + if not isinstance(sample_rows_in_table_info, int): + raise TypeError("sample_rows_in_table_info must be an integer") + + self._sample_rows_in_table_info = sample_rows_in_table_info + self._indexes_in_table_info = indexes_in_table_info + + self._custom_table_info = custom_table_info + if self._custom_table_info: + if not isinstance(self._custom_table_info, dict): + raise TypeError( + "table_info must be a dictionary with table names as keys and the " + "desired table info as values" + ) + # only keep the tables that are also present in the database + intersection = set(self._custom_table_info).intersection(self._all_tables) + self._custom_table_info = dict( + (table, self._custom_table_info[table]) + for table in self._custom_table_info + if table in intersection + ) + + self._max_string_length = max_string_length + + self._metadata = metadata or MetaData() + # including view support if view_support = true + self._metadata.reflect( + views=view_support, + bind=self._engine, + only=list(self._usable_tables), + schema=self._schema, + ) + + @classmethod + def from_uri( + cls, database_uri: str, engine_args: Optional[dict] = None, **kwargs: Any + ) -> SQLDatabase: + """Construct a SQLAlchemy engine from URI.""" + _engine_args = engine_args or {} + return cls(create_engine(database_uri, **_engine_args), **kwargs) + + @classmethod + def from_databricks( + cls, + catalog: str, + schema: str, + host: Optional[str] = None, + api_token: Optional[str] = None, + warehouse_id: Optional[str] = None, + cluster_id: Optional[str] = None, + engine_args: Optional[dict] = None, + **kwargs: Any, + ) -> SQLDatabase: + """ + Class method to create an SQLDatabase instance from a Databricks connection. + This method requires the 'databricks-sql-connector' package. If not installed, + it can be added using `pip install databricks-sql-connector`. + + Args: + catalog (str): The catalog name in the Databricks database. + schema (str): The schema name in the catalog. + host (Optional[str]): The Databricks workspace hostname, excluding + 'https://' part. If not provided, it attempts to fetch from the + environment variable 'DATABRICKS_HOST'. If still unavailable and if + running in a Databricks notebook, it defaults to the current workspace + hostname. Defaults to None. + api_token (Optional[str]): The Databricks personal access token for + accessing the Databricks SQL warehouse or the cluster. If not provided, + it attempts to fetch from 'DATABRICKS_TOKEN'. If still unavailable + and running in a Databricks notebook, a temporary token for the current + user is generated. Defaults to None. + warehouse_id (Optional[str]): The warehouse ID in the Databricks SQL. If + provided, the method configures the connection to use this warehouse. + Cannot be used with 'cluster_id'. Defaults to None. + cluster_id (Optional[str]): The cluster ID in the Databricks Runtime. If + provided, the method configures the connection to use this cluster. + Cannot be used with 'warehouse_id'. If running in a Databricks notebook + and both 'warehouse_id' and 'cluster_id' are None, it uses the ID of the + cluster the notebook is attached to. Defaults to None. + engine_args (Optional[dict]): The arguments to be used when connecting + Databricks. Defaults to None. + **kwargs (Any): Additional keyword arguments for the `from_uri` method. + + Returns: + SQLDatabase: An instance of SQLDatabase configured with the provided + Databricks connection details. + + Raises: + ValueError: If 'databricks-sql-connector' is not found, or if both + 'warehouse_id' and 'cluster_id' are provided, or if neither + 'warehouse_id' nor 'cluster_id' are provided and it's not executing + inside a Databricks notebook. + """ + try: + from databricks import sql # noqa: F401 + except ImportError: + raise ValueError( + "databricks-sql-connector package not found, please install with" + " `pip install databricks-sql-connector`" + ) + context = None + try: + from dbruntime.databricks_repl_context import get_context + + context = get_context() + except ImportError: + pass + + default_host = context.browserHostName if context else None + if host is None: + host = get_from_env("host", "DATABRICKS_HOST", default_host) + + default_api_token = context.apiToken if context else None + if api_token is None: + api_token = get_from_env("api_token", "DATABRICKS_TOKEN", default_api_token) + + if warehouse_id is None and cluster_id is None: + if context: + cluster_id = context.clusterId + else: + raise ValueError( + "Need to provide either 'warehouse_id' or 'cluster_id'." + ) + + if warehouse_id and cluster_id: + raise ValueError("Can't have both 'warehouse_id' or 'cluster_id'.") + + if warehouse_id: + http_path = f"/sql/1.0/warehouses/{warehouse_id}" + else: + http_path = f"/sql/protocolv1/o/0/{cluster_id}" + + uri = ( + f"databricks://token:{api_token}@{host}?" + f"http_path={http_path}&catalog={catalog}&schema={schema}" + ) + return cls.from_uri(database_uri=uri, engine_args=engine_args, **kwargs) + + @classmethod + def from_cnosdb( + cls, + url: str = "127.0.0.1:8902", + user: str = "root", + password: str = "", + tenant: str = "cnosdb", + database: str = "public", + ) -> SQLDatabase: + """ + Class method to create an SQLDatabase instance from a CnosDB connection. + This method requires the 'cnos-connector' package. If not installed, it + can be added using `pip install cnos-connector`. + + Args: + url (str): The HTTP connection host name and port number of the CnosDB + service, excluding "http://" or "https://", with a default value + of "127.0.0.1:8902". + user (str): The username used to connect to the CnosDB service, with a + default value of "root". + password (str): The password of the user connecting to the CnosDB service, + with a default value of "". + tenant (str): The name of the tenant used to connect to the CnosDB service, + with a default value of "cnosdb". + database (str): The name of the database in the CnosDB tenant. + + Returns: + SQLDatabase: An instance of SQLDatabase configured with the provided + CnosDB connection details. + """ + try: + from cnosdb_connector import make_cnosdb_langchain_uri + + uri = make_cnosdb_langchain_uri(url, user, password, tenant, database) + return cls.from_uri(database_uri=uri) + except ImportError: + raise ValueError( + "cnos-connector package not found, please install with" + " `pip install cnos-connector`" + ) + + @property + def dialect(self) -> str: + """Return string representation of dialect to use.""" + return self._engine.dialect.name + + def get_usable_table_names(self) -> Iterable[str]: + """Get names of tables available.""" + if self._include_tables: + return sorted(self._include_tables) + return sorted(self._all_tables - self._ignore_tables) + + def get_table_names(self) -> Iterable[str]: + """Get names of tables available.""" + warnings.warn( + "This method is deprecated - please use `get_usable_table_names`." + ) + return self.get_usable_table_names() + + @property + def table_info(self) -> str: + """Information about all tables in the database.""" + return self.get_table_info() + + def get_table_info(self, table_names: Optional[List[str]] = None) -> str: + """Get information about specified tables. + + Follows best practices as specified in: Rajkumar et al, 2022 + (https://arxiv.org/abs/2204.00498) + + If `sample_rows_in_table_info`, the specified number of sample rows will be + appended to each table description. This can increase performance as + demonstrated in the paper. + """ + all_table_names = self.get_usable_table_names() + if table_names is not None: + missing_tables = set(table_names).difference(all_table_names) + if missing_tables: + raise ValueError(f"table_names {missing_tables} not found in database") + all_table_names = table_names + + meta_tables = [ + tbl + for tbl in self._metadata.sorted_tables + if tbl.name in set(all_table_names) + and not (self.dialect == "sqlite" and tbl.name.startswith("sqlite_")) + ] + + tables = [] + for table in meta_tables: + if self._custom_table_info and table.name in self._custom_table_info: + tables.append(self._custom_table_info[table.name]) + continue + + # Ignore JSON datatyped columns + for k, v in table.columns.items(): + if type(v.type) is NullType: + table._columns.remove(v) + + # add create table command + create_table = str(CreateTable(table).compile(self._engine)) + table_info = f"{create_table.rstrip()}" + has_extra_info = ( + self._indexes_in_table_info or self._sample_rows_in_table_info + ) + if has_extra_info: + table_info += "\n\n/*" + if self._indexes_in_table_info: + table_info += f"\n{self._get_table_indexes(table)}\n" + if self._sample_rows_in_table_info: + table_info += f"\n{self._get_sample_rows(table)}\n" + if has_extra_info: + table_info += "*/" + tables.append(table_info) + tables.sort() + final_str = "\n\n".join(tables) + return final_str + + def _get_table_indexes(self, table: Table) -> str: + indexes = self._inspector.get_indexes(table.name) + indexes_formatted = "\n".join(map(_format_index, indexes)) + return f"Table Indexes:\n{indexes_formatted}" + + def _get_sample_rows(self, table: Table) -> str: + # build the select command + command = select(table).limit(self._sample_rows_in_table_info) + + # save the columns in string format + columns_str = "\t".join([col.name for col in table.columns]) + + try: + # get the sample rows + with self._engine.connect() as connection: + sample_rows_result = connection.execute(command) # type: ignore + # shorten values in the sample rows + sample_rows = list( + map(lambda ls: [str(i)[:100] for i in ls], sample_rows_result) + ) + + # save the sample rows in string format + sample_rows_str = "\n".join(["\t".join(row) for row in sample_rows]) + + # in some dialects when there are no rows in the table a + # 'ProgrammingError' is returned + except ProgrammingError: + sample_rows_str = "" + + return ( + f"{self._sample_rows_in_table_info} rows from {table.name} table:\n" + f"{columns_str}\n" + f"{sample_rows_str}" + ) + + def _execute( + self, + command: str, + fetch: Union[Literal["all"], Literal["one"]] = "all", + ) -> Sequence[Dict[str, Any]]: + """ + Executes SQL command through underlying engine. + + If the statement returns no rows, an empty list is returned. + """ + with self._engine.begin() as connection: + if self._schema is not None: + if self.dialect == "snowflake": + connection.exec_driver_sql( + "ALTER SESSION SET search_path = %s", (self._schema,) + ) + elif self.dialect == "bigquery": + connection.exec_driver_sql("SET @@dataset_id=?", (self._schema,)) + elif self.dialect == "mssql": + pass + elif self.dialect == "trino": + connection.exec_driver_sql("USE ?", (self._schema,)) + elif self.dialect == "duckdb": + # Unclear which parameterized argument syntax duckdb supports. + # The docs for the duckdb client say they support multiple, + # but `duckdb_engine` seemed to struggle with all of them: + # https://github.com/Mause/duckdb_engine/issues/796 + connection.exec_driver_sql(f"SET search_path TO {self._schema}") + elif self.dialect == "oracle": + connection.exec_driver_sql( + f"ALTER SESSION SET CURRENT_SCHEMA = {self._schema}" + ) + else: # postgresql and other compatible dialects + connection.exec_driver_sql("SET search_path TO %s", (self._schema,)) + cursor = connection.execute(text(command)) + if cursor.returns_rows: + if fetch == "all": + result = [x._asdict() for x in cursor.fetchall()] + elif fetch == "one": + first_result = cursor.fetchone() + result = [] if first_result is None else [first_result._asdict()] + else: + raise ValueError("Fetch parameter must be either 'one' or 'all'") + return result + return [] + + def run( + self, + command: str, + fetch: Union[Literal["all"], Literal["one"]] = "all", + ) -> str: + """Execute a SQL command and return a string representing the results. + + If the statement returns rows, a string of the results is returned. + If the statement returns no rows, an empty string is returned. + """ + result = self._execute(command, fetch) + # Convert columns values to string to avoid issues with sqlalchemy + # truncating text + res = [ + tuple(truncate_word(c, length=self._max_string_length) for c in r.values()) + for r in result + ] + if not res: + return "" + else: + return str(res) + + def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str: + """Get information about specified tables. + + Follows best practices as specified in: Rajkumar et al, 2022 + (https://arxiv.org/abs/2204.00498) + + If `sample_rows_in_table_info`, the specified number of sample rows will be + appended to each table description. This can increase performance as + demonstrated in the paper. + """ + try: + return self.get_table_info(table_names) + except ValueError as e: + """Format the error message""" + return f"Error: {e}" + + def run_no_throw( + self, + command: str, + fetch: Union[Literal["all"], Literal["one"]] = "all", + ) -> str: + """Execute a SQL command and return a string representing the results. + + If the statement returns rows, a string of the results is returned. + If the statement returns no rows, an empty string is returned. + + If the statement throws an error, the error message is returned. + """ + try: + return self.run(command, fetch) + except SQLAlchemyError as e: + """Format the error message""" + return f"Error: {e}" diff --git a/src/aios_kernel/sql_database_function.py b/src/aios_kernel/sql_database_function.py new file mode 100644 index 0000000..2e208e6 --- /dev/null +++ b/src/aios_kernel/sql_database_function.py @@ -0,0 +1,112 @@ +from datetime import timedelta, datetime +from typing import Dict + +from cachetools import TLRUCache, cached + +from aios_kernel.ai_function import AIFunction +from aios_kernel.sql_database import SQLDatabase, get_from_env + + +def _my_ttu(_key, _value, now): + return now + timedelta(seconds=600) + + +database_cache = TLRUCache(ttu=_my_ttu, maxsize=10000, timer=datetime.now) + + +@cached(cache=database_cache) +def get_database(uri: str) -> SQLDatabase: + return SQLDatabase.from_uri(uri) + + +class GetTableInfosFunction(AIFunction): + def __init__(self): + super().__init__() + self.name = "get_table_infos" + self.description = "Get table informations in the database" + + def get_name(self) -> str: + return self.name + + def get_description(self) -> str: + return self.description + + def get_parameters(self) -> Dict: + return { + "type": "object", + "properties": { + "database_url": {"type": "string", "description": "Database URL,Can be set to None"}, + } + } + + async def execute(self, **kwargs) -> str: + database_url: str = kwargs.get("database_url") + if (database_url is None + or database_url.strip() == "" + or database_url.strip().lower() == "none" + or database_url.strip().lower() == "null"): + database_url = get_from_env(key="database url", env_key="DATABASE_URL") + if database_url is None: + return "error: database_url is None" + database = get_database(database_url) + tables = database.get_usable_table_names() + table_infos = database.get_table_info(tables) + return table_infos + + def is_local(self) -> bool: + return True + + def is_in_zone(self) -> bool: + return True + + def is_ready_only(self) -> bool: + return False + + +class ExecuteSqlFunction(AIFunction): + def __init__(self): + super().__init__() + self.name = "execute_sql" + self.description = """ + Input to this function is a detailed and correct SQL query, output is a result from the database. + If the query is not correct, an error message will be returned. + If an error is returned, rewrite the query, check the query, and try again. + """ + + def get_name(self) -> str: + return self.name + + def get_description(self) -> str: + return self.description + + def get_parameters(self) -> Dict: + return { + "type": "object", + "properties": { + "database_url": {"type": "string", "description": "Database URL,Can be set to None"}, + "sql": {"type": "string", "description": "SQL to execute"} + } + } + + async def execute(self, **kwargs) -> str: + database_url = kwargs.get("database_url") + if (database_url is None + or database_url.strip() == "" + or database_url.strip().lower() == "none" + or database_url.strip().lower() == "null"): + database_url = get_from_env(key="database url", env_key="DATABASE_URL") + if database_url is None: + return "error: database_url is None" + sql = kwargs.get("sql") + + database = get_database(database_url) + return database.run_no_throw(sql) + + def is_local(self) -> bool: + return True + + def is_in_zone(self) -> bool: + return True + + def is_ready_only(self) -> bool: + return False diff --git a/src/component/llama_node/local_llama_compute_node.py b/src/component/llama_node/local_llama_compute_node.py index 0a22369..efb9712 100644 --- a/src/component/llama_node/local_llama_compute_node.py +++ b/src/component/llama_node/local_llama_compute_node.py @@ -1,9 +1,6 @@ -import json import logging import requests -from typing import Optional, List -from pydantic import BaseModel from aios import ComputeTask,Queue_ComputeNode, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType,AIStorage,UserConfig diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 0ac7a81..1a13476 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -28,6 +28,8 @@ sys.path.append(directory + '/../../') import proxy from aios import * +import local_compute_node_builder +from component.llama_node.local_llama_compute_node import LocalLlama_ComputeNode sys.path.append(directory + '/../../component/') @@ -402,41 +404,34 @@ class AIOS_Shell: async def handle_node_commands(self, args): show_text = FormattedText([("class:title", "sub command not support!\n" - "/node add llama $model_name $url\n" - "/node rm llama $model_name $url\n" + "/node add $model_name $url\n" + "/node create\n" + "/node rm $model_name $url\n" "/node list\n")]) if len(args) < 1: return show_text sub_cmd = args[0] - if sub_cmd == "add": - if len(args) < 2: - return show_text - if args[1] == "llama": - if len(args) < 4: - return show_text - - model_name = args[2] - url = args[3] - ComputeNodeConfig.get_instance().add_node("llama", url, model_name) - ComputeNodeConfig.get_instance().save() - node = LocalLlama_ComputeNode(url, model_name) - node.start() - ComputeKernel.get_instance().add_compute_node(node) - else: + if sub_cmd == "create": + await local_compute_node_builder.build(session, shell_style) + elif sub_cmd == "add": + if len(args) < 3: return show_text + + model_name = args[1] + url = args[2] + ComputeNodeConfig.get_instance().add_node("llama", url, model_name) + ComputeNodeConfig.get_instance().save() + node = LocalLlama_ComputeNode(url, model_name) + node.start() + ComputeKernel.get_instance().add_compute_node(node) elif sub_cmd == "rm": - if len(args) < 2: - return show_text - if args[1] == "llama": - if len(args) < 4: - return show_text - - model_name = args[2] - url = args[3] - ComputeNodeConfig.get_instance().remove_node("llama", url, model_name) - ComputeNodeConfig.get_instance().save() - else: + if len(args) < 3: return show_text + + model_name = args[1] + url = args[2] + ComputeNodeConfig.get_instance().remove_node("llama", url, model_name) + ComputeNodeConfig.get_instance().save() elif sub_cmd == "list": print_formatted_text(ComputeNodeConfig.get_instance().list()) @@ -785,8 +780,9 @@ async def main(): '/set_config $key', '/enable $feature', '/disable $feature', - '/node add llama $model_name $url', - '/node rm llama $model_name $url', + '/node add $model_name $url', + '/node create', + '/node rm $model_name $url', '/node list', '/show', '/exit', diff --git a/src/service/aios_shell/local_compute_node_builder/__init__.py b/src/service/aios_shell/local_compute_node_builder/__init__.py new file mode 100644 index 0000000..a04f822 --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/__init__.py @@ -0,0 +1,38 @@ +import os +from prompt_toolkit import HTML, PromptSession, print_formatted_text +from prompt_toolkit.styles import Style +from aios.storage.storage import AIStorage +from service.aios_shell.local_compute_node_builder.local_llama_node_builder import LocalLlamaNodeBuilder +from .local_compute_node_builder import BuilderState + +async def build(prompt_session: PromptSession, shell_style: Style) -> str or None: + # model_type = await prompt_session.prompt_async(f"Please select the node server type (default: llama.cpp):", style = shell_style) + + model_type = 'llama.cpp' + + download_dir = AIStorage.get_instance().get_download_dir() + if not os.path.exists(download_dir): + os.mkdir(download_dir) + + state = BuilderState(prompt_session, shell_style) + + match model_type: + case 'llama.cpp': + builder = LocalLlamaNodeBuilder(state) + + while True: + param = builder.next_parameter() + if param is None: + return None + + if state.last_result_prompt or param.desc: + print_formatted_text(f"{state.last_result_prompt}{param.desc}", style = state.shell_style) + value = await state.prompt_session.prompt_async(f"{param.prompt}:", style = state.shell_style) + if value: + value = value.strip() + + state.params[param.name] = value + url = await param.applier.apply(state, param.name, value) + + if url is not None: + return url \ No newline at end of file diff --git a/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py b/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py new file mode 100644 index 0000000..33e1e36 --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/local_compute_node_builder.py @@ -0,0 +1,40 @@ +from abc import abstractmethod + +from prompt_toolkit import PromptSession +from prompt_toolkit.styles import Style + +class BuilderState: + def __init__(self, prompt_session: PromptSession, shell_style: Style): + self.prompt_session = prompt_session + self.shell_style = shell_style + self.next_step = 0 + self.last_result_prompt = "" + self.params = {} + +# class ApplyResult: +# def __init__(self, next_step: any, url: str or None = None, result_prompt: str or None = None) -> None: +# self.next_step = next_step +# self.url = url +# self.result_prompt = result_prompt + + +class ParameterApplier: + @abstractmethod + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + pass + +class BuildParameter: + def __init__(self, name: str, applier: ParameterApplier, prompt: str or None = None, desc: str or None = None, default_value: str or None = None): + self.name = name + self.prompt = prompt + self.desc = desc + self.default_value = default_value + self.applier = applier + +class LocalComputeNodeBuilder: + def __init__(self, state: BuilderState) -> None: + self.state = state + + @abstractmethod + def next_parameter(self) -> BuildParameter or None: + pass \ No newline at end of file diff --git a/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py b/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py new file mode 100644 index 0000000..d9bdcde --- /dev/null +++ b/src/service/aios_shell/local_compute_node_builder/local_llama_node_builder.py @@ -0,0 +1,254 @@ +import os +import random +import subprocess +import requests + +from prompt_toolkit import print_formatted_text +from prompt_toolkit.shortcuts import ProgressBar +from prompt_toolkit.formatted_text import FormattedText + +from aios.storage.storage import AIStorage +from aios import ComputeKernel +from component.llama_node.local_llama_compute_node import LocalLlama_ComputeNode +from service.aios_shell.compute_node_config import ComputeNodeConfig +from .local_compute_node_builder import BuildParameter, BuilderState, LocalComputeNodeBuilder, ParameterApplier + +class BuildParameterModelPath: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value: + if os.path.exists(value): + state.next_step += 2 + else: + print_formatted_text(FormattedText([("class:error", f"Model not exist at {value}")]), style = state.shell_style) + else: + state.next_step += 1 + + +class BuildParameterModelUrl: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value is None: + value = "1" + + url = value + recommend = _recommend_model_urls.get(value) + if recommend: + url = recommend["url"] + + save_path = f"{AIStorage.get_instance().get_download_dir()}/{url.split('/').pop()}" + + print_formatted_text(FormattedText([("class:prompt", f"Will save the model to {save_path}:\n")]), style = state.shell_style) + + try: + # get file size + response = requests.head(url) + file_size = int(response.headers.get('content-length', 0)) + + # start download + response = requests.get(url, stream=True) + + if response.status_code == 200: + with open(save_path, 'wb') as f, ProgressBar() as pb: + for data in pb(response.iter_content(1024), total = (file_size + 1023) // 1024): + f.write(data) + + print_formatted_text(FormattedText([("class:prompt", f"Download model success, save at: {save_path}\n")]), style = state.shell_style) + + state.params["model_path"] = save_path + state.next_step += 1 + else: + print_formatted_text(FormattedText([("class:error", f"Download model failed, error: {response.status_code}\nYou can retry it or select another one.")]), style = state.shell_style) + + except Exception as e: + print_formatted_text(FormattedText([("class:error", f"Download model failed: {e}\nYou can retry it or select another one.")]), style = state.shell_style) + +class ParameterNodeNameApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or os.path.basename(state.params["model_path"]) + state.params["node_name"] = value + state.next_step += 1 + +class ParameterPortApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + if value is None or value == "0": + value = str(random.randint(10000, 60000)) + + state.params["port"] = value + state.next_step += 1 + +class ParameterNGpuLayersApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "83" + state.params["n_gpu_layers"] = value + state.next_step += 1 + +class ParameterNCtxApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "4096" + state.params["n_ctx"] = value + state.next_step += 1 + +class ParameterChatFormatApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + value = value or "llama-2" + state.params["chat_format"] = value + state.next_step += 1 + +class ParameterExternParamsApplier: + async def apply(self, state: BuilderState, name: str, value: str or None = None) -> str or None: + extern_params = value + docker_image = "" + gpu_options = [] + state.next_step += 1 + + if state.params["n_gpu_layers"] == "0": + docker_image = "ghcr.io/abetlen/llama-cpp-python:latest" + else: + gpu_options = ["--gpus", "all"] + llama_cpp_python_repo_url = "https://github.com/abetlen/llama-cpp-python.git" + download_path = AIStorage.get_instance().get_download_dir() + llama_cpp_python_path = download_path + "/llama-cpp-python" + + # update the `llama-cpp-python` + retry = True + while retry: + retry = False + result = None + if os.path.exists(llama_cpp_python_path): + result = subprocess.run(['git', 'pull'], cwd = llama_cpp_python_path, stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + else: + result = subprocess.run(['git', 'clone', llama_cpp_python_repo_url, llama_cpp_python_path], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + print_formatted_text(FormattedText([("class:warn", result.stderr)]), style = state.shell_style) + while True: + sel = await state.prompt_session.prompt_async(f"Update 'llama-cpp-python' failed, you can press 'r' to retry, or 'c' to continue with the current version.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'c': + break + else: + pass # Select again + else: + break + + # build the image + docker_image = 'llama-cpp-python-cuda' + retry = True + while retry: + retry = False + result = subprocess.run(['docker', 'rmi', docker_image], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + result = subprocess.run(['docker', 'build', '-t', docker_image, f"{llama_cpp_python_path}/docker/cuda_simple/"], stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + print_formatted_text(FormattedText([("class:warn", result.stderr)]), style = state.shell_style) + while True: + sel = await state.prompt_session.prompt_async(f"Build the image failed, you can press 'r' to retry, or 'c' to continue with the current version.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'c': + break + else: + pass # Select again + else: + break + + retry = True + while retry: + retry = False + run_options = ['docker', 'run', '-d'] + + if gpu_options: + run_options.extend(gpu_options) + + run_options.extend([ + '-p', f"{state.params['port']}:8000", + '-v', f"{os.path.dirname(state.params['model_path'])}:/models", '-e', f"MODEL=/models/{os.path.basename(state.params['model_path'])}", + 'llama-cpp-python-cuda', + 'python3', '-m', 'llama_cpp.server', + '--n_gpu_layers', state.params["n_gpu_layers"], + '--n_ctx', state.params["n_ctx"], + '--chat_format', state.params["chat_format"], + ]) + + if extern_params: + run_options.extend(extern_params.split(' ')) + + print_formatted_text(FormattedText([("class:prompt", f"Will start service with: {' '.join(run_options)}")]), style = state.shell_style) + + result = subprocess.run(run_options, stdout = subprocess.PIPE, stderr = subprocess.PIPE, text = True) + + if result.stderr: + print_formatted_text(FormattedText([("class:warn", result.stderr)]), style = state.shell_style) + while True: + sel = await state.prompt_session.prompt_async(f"Start the node service failed, you can press 'r' to retry, or 'a' to abort.", style = state.shell_style) + if sel == 'r': + retry = True + break + elif sel == 'a': + break + else: + pass # Select again + else: + local_url = f'http://localhost:{state.params["port"]}' + foreign_url = 'http://{your-host-address}:' + state.params["port"] + model_name = state.params['node_name'] + + ComputeNodeConfig.get_instance().add_node("llama", local_url, model_name) + ComputeNodeConfig.get_instance().save() + node = LocalLlama_ComputeNode(local_url, model_name) + node.start() + ComputeKernel.get_instance().add_compute_node(node) + + print_formatted_text(FormattedText([( + "class:prompt", +f""" +Congratulations! The node ({model_name}) service successed. +You can access it with follow url: +{local_url} +And 'http://{foreign_url}' in other computers. +Now you can refer it in agents as `llm_model_name={model_name}` +""" + )]), style = state.shell_style) + break + +_recommend_model_urls = { + "1": { + "model": "Llama-2-70B-Chat-GGUF", + "url": "https://huggingface.co/TheBloke/Llama-2-70B-chat-GGUF/resolve/main/llama-2-70b-chat.Q4_0.gguf" + }, + "2": { + "model": "Llama-2-13B-Chat-GGUF", + "url": "https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF/resolve/main/llama-2-13b-chat.Q4_0.gguf" + }, + "3": { + "model": "Llama-2-7B-Chat-GGUF", + "url": "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q4_K_M.gguf" + }, +} + +_recommend_model_url_table_str = "" +for i in range(1, 999): + id = str(i) + info = _recommend_model_urls.get(id) + if info: + _recommend_model_url_table_str += f"\n\t{id}\t{info['model']}\t{info['url']}" + else: + break + +_params = [ + BuildParameter("model_path", BuildParameterModelPath(), "Please input the model file path (Press 'Enter' if you need to download it)"), + BuildParameter("model_url", BuildParameterModelUrl(), "Please input (default: Llama-2-70B-chat)", f"Now you need input the url to download the model, or you can input the 'ID' in the follow table to select one:\n\tID\tmodel\t\turl{_recommend_model_url_table_str}"), + BuildParameter("node_name", ParameterNodeNameApplier(), "Please input name for your node, and you can set it in 'llm_model_name' of 'agent.toml' (default: the name of the model file)"), + BuildParameter("port", ParameterPortApplier(), "Please input the port which the node server will listen on (default: random)"), + BuildParameter("n_gpu_layers", ParameterNGpuLayersApplier(), "Please input layers offload to GPU (<=83 for Llama, 0 for CPU only, default: 83)"), + BuildParameter("n_ctx", ParameterNCtxApplier(), "Please input the content limit (default: 4096)"), + BuildParameter("chat_format", ParameterChatFormatApplier(), "Please input the chat format (default: llama-2)"), + BuildParameter("extern_params", ParameterExternParamsApplier(), "Please input other parameters refer to 'llama-cpp-python'(https://github.com/abetlen/llama-cpp-python), press 'Enter' to ignore it"), +] + +class LocalLlamaNodeBuilder(LocalComputeNodeBuilder): + def next_parameter(self) -> BuildParameter or None: + if self.state.next_step < len(_params): + return _params[self.state.next_step]