Merge pull request #107 from streetycat/MVP
Compute Node Installation Wizard
This commit is contained in:
Liu Zhicong
GitHub
@@ -217,6 +217,13 @@ class AIStorage:
|
||||
"""
|
||||
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)}")
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
if sub_cmd == "create":
|
||||
await local_compute_node_builder.build(session, shell_style)
|
||||
elif sub_cmd == "add":
|
||||
if len(args) < 3:
|
||||
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:
|
||||
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:
|
||||
if len(args) < 3:
|
||||
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:
|
||||
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',
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user