KlonetAI: An AI agent for intelligent interaction with Klonet.
import zhipuai
from key import ZhiPuAI_API_Key
from transformers.tools import Agent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT


class ChatGLMAgent(Agent):
    """
    Agent that uses ChatGLM Inference SDK to generate code.

    Args:
        model_name (`str`): The name of the model to be used, could be "chatglm_pro", "chatglm_std", "chatglm_lite",
            or "chatglm_lite_32k".
        key (`str`, *optional*): The API key obtained from ZhiPuAI Dashboard.
        chat_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `chat_prompt_template.txt` in this repo in this case.
        run_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `run_prompt_template.txt` in this repo in this case.
        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
            one of the default tools, that default tool will be overridden.

    Example:

    ```py
    agent = ChatGLMAgent("chatglm_std")
    agent.run("Create a star topology with 5 host nodes and interconnect them with an OVS switch.")
    ```
    """

    def __init__(self, model_name, key="", chat_prompt_template=None, run_prompt_template=None, additional_tools=None):
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )
        self.model_name = model_name
        zhipuai.api_key = key or ZhiPuAI_API_Key
        self._mode = "agent"

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def generate_one(self, prompt, stop):
        response = zhipuai.model_api.invoke(
            model=self.model_name,
            prompt=[{"role": "user", "content": prompt}],
            temperature=0,
        )

        if response["code"] != 200:
            raise ValueError(f"Error {response.status_code}: {response}")

        result = response["data"]["choices"][0]["content"]\
            .strip('"').replace('\\n', '\n').replace('\\"', '\"')

        # Inference API returns the stop sequence
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/chatglm.py
Lizonghang/KlonetAI
import dashscope
from key import DashScope_API_Key
from transformers.tools import Agent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT


class DashScopeAgent(Agent):
    """
    Agent that uses Aliyun DashScope Inference SDK to generate code.

    Args:
        model_name (`str`): The name of the model to be used, could be "ernie-bot" or "ernie-bot-turbo".
        key (`str`, *optional*): The API key obtained from Aliyun Dashboard.
        chat_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `chat_prompt_template.txt` in this repo in this case.
        run_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `run_prompt_template.txt` in this repo in this case.
        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
            one of the default tools, that default tool will be overridden.

    Example:

    ```py
    agent = DashScopeAgent("ernie-bot-turbo")
    agent.run("Create a star topology with 5 host nodes and interconnect them with an OVS switch.")
    ```
    """

    def __init__(self, model_name, key="", chat_prompt_template=None, run_prompt_template=None, additional_tools=None):
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )
        self.model_name = model_name
        dashscope.api_key = key or DashScope_API_Key
        self._gen = dashscope.Generation()
        self._mode = "agent"

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def generate_one(self, prompt, stop):
        response = self._gen.call(
            model=self.model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
        )

        if response["status_code"] != 200:
            raise ValueError(f"Error {response.status_code}: {response}")

        result = response["output"]["text"]
        # Inference API returns the stop sequence
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/dashscope.py
Lizonghang/KlonetAI
import erniebot
from key import AI_Studio_APT_Key
from transformers.tools import Agent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT


class ErnieBotAgent(Agent):
    """
    Agent that uses Baidu Ernie Bot Inference SDK to generate code.

    Args:
        model_name (`str`): The name of the model to be used, could be "ernie-bot" or "ernie-bot-turbo".
        key (`str`, *optional*): The API key obtained from Baidu Ernie Dashboard.
        chat_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `chat_prompt_template.txt` in this repo in this case.
        run_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `run_prompt_template.txt` in this repo in this case.
        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
            one of the default tools, that default tool will be overridden.

    Example:

    ```py
    agent = ErnieBotAgent("ernie-bot-turbo")
    agent.run("Create a star topology with 5 host nodes and interconnect them with an OVS switch.")
    ```
    """

    def __init__(self, model_name, key="", chat_prompt_template=None, run_prompt_template=None, additional_tools=None):
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )
        self.model_name = model_name
        erniebot.api_type = "aistudio"
        erniebot.access_token = key or AI_Studio_APT_Key
        self._mode = "agent"

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def generate_one(self, prompt, stop):
        response = erniebot.ChatCompletion.create(
            model=self.model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
        )

        if response["code"] != 200:
            raise ValueError(f"Error {response.status_code}: {response}")

        result = response["result"]
        # Inference API returns the stop sequence
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/erniebot.py
Lizonghang/KlonetAI
from transformers.tools import HfAgent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT


class KAIHfAgent(HfAgent):

    def __init__(self, *args, **kwargs):
        self._mode = "agent"
        super().__init__(*args, **kwargs)

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/huggingface.py
Lizonghang/KlonetAI
from transformers.tools import OpenAiAgent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT


class KAIOpenAIAgent(OpenAiAgent):

    def __init__(self, *args, **kwargs):
        self._mode = "agent"
        super().__init__(*args, **kwargs)

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/openai.py
Lizonghang/KlonetAI
import ssl
import hmac
import json
import base64
import hashlib
import websocket
import _thread as thread
from urllib.parse import urlencode, urlparse
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
from key import SparkDesk_AppId, SparkDesk_AppSecret, SparkDesk_API_Key
from transformers.tools import Agent
from transformers.tools.prompts import CHAT_MESSAGE_PROMPT

websocket.enableTrace(False)
result = ""


def on_error(ws, error):
    print("[WebSocket Error]", error)


def on_close(ws, args1, args2):
    print("[WebSocket Closed]")


def on_open(ws):
    thread.start_new_thread(run, (ws,))


def on_message(ws, message):
    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        print(f'Error {code}: {data}')
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        global result
        result += content
        if status == 2:
            ws.close()


def run(ws, *args):
    ws.send(json.dumps(gen_params(ws.prompt)))


def gen_params(prompt):
    return {
        "header": {
            "app_id": SparkDesk_AppId,
            "uid": "klonetai-user"
        },
        "parameter": {
            "chat": {
                "domain": "generalv2",
                "temperature": 0.1,
                "max_tokens": 2048,
            }
        },
        "payload": {
            "message": {
                "text": [
                    {"role": "user", "content": prompt},
                ]
            }
        }
    }


class Ws_Param(object):
    def __init__(self, APPID, APIKey, APISecret, Spark_URL):
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(Spark_URL).netloc
        self.path = urlparse(Spark_URL).path
        self.Spark_url = Spark_URL

    def create_url(self):
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        signature_sha = hmac.new(
            self.APISecret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'

        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        url = self.Spark_url + '?' + urlencode(v)
        return url


class SparkDesk(Agent):
    """
    Agent that uses Xunfei SparkDesk REST API to generate code.

    Args:
        key (`str`, *optional*): The API key obtained from Xunfei SparkDesk Dashboard.
        appid (`str`, *optional*): The App Id obtained from Xunfei SparkDesk Dashboard.
        appsecret (`str`, *optional*): The App Secret obtained from Xunfei SparkDesk Dashboard.
        chat_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `chat_prompt_template.txt` in this repo in this case.
        run_prompt_template (`str`, *optional*):
            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
            `run_prompt_template.txt` in this repo in this case.
        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
            one of the default tools, that default tool will be overridden.

    Example:

    ```py
    agent = SparkDesk()
    agent.run("Create a star topology with 5 host nodes and interconnect them with an OVS switch.")
    ```
    """

    def __init__(self, key="", appid="", appsecret="", chat_prompt_template=None,
                 run_prompt_template=None, additional_tools=None):
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )
        self.model_name = "spark-v2"
        if key: SparkDesk_API_Key = key
        if appid: SparkDesk_AppId = appid
        if appsecret: SparkDesk_AppSecret = appsecret
        self._mode = "agent"

    @property
    def mode(self):
        return self._mode

    def set_mode(self, mode):
        self._mode = mode

    def generate_one(self, prompt, stop):
        global result
        result = ""

        wsParam = Ws_Param(
            APPID=SparkDesk_AppId,
            APIKey=SparkDesk_API_Key,
            APISecret=SparkDesk_AppSecret,
            Spark_URL="ws://spark-api.xf-yun.com/v2.1/chat"
        )
        ws = websocket.WebSocketApp(
            wsParam.create_url(),
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
            on_open=on_open
        )
        ws.prompt = prompt
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

        # Inference API returns the stop sequence
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        if not self._mode == "chat":
            return super().chat(task, return_code=return_code, remote=remote, **kwargs)

        if self.chat_history is None:
            description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
            prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
        else:
            prompt = self.chat_history
        prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", f"[Chat only] {task}")
        response = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + response.strip() + "\n"
        self.log(response + "\n")

agent/sparkdesk.py
Lizonghang/KlonetAI
import io
import os
import os.path
import requests.exceptions
import tool
import tutorial
import panel as pn
from datetime import datetime
from panel.widgets import SpeechToText
from contextlib import redirect_stdout
from tool.klonet import kai

tools = [obj() for obj in tool.base]        # 119 tokens
tools += [obj() for obj in tool.klonet]     # 4922 tokens
tools += [obj() for obj in tool.topo]       # 541 tokens
tools += [obj() for obj in tool.gpt]        # 184 tokens
tools += [obj() for obj in tutorial.base]   # 121 tokens
tools += [obj() for obj in tutorial.mxnet]

chat_time = None
save_dir = "output"
os.makedirs(save_dir, exist_ok=True)

pn.extension("floatpanel", notifications=True)

# Login to Klonet backend.
project_name_input = pn.widgets.TextInput(name="Project:", value="")
user_name_input = pn.widgets.TextInput(name="User:", value="")
host_ip_input = pn.widgets.PasswordInput(name="Klonet Backend Host IP:", value="")
port_input = pn.widgets.PasswordInput(name="Klonet Backend Port:", value="")


def click_login_klonet_button(event):
    project_name = project_name_input.value
    user_name = user_name_input.value
    host_ip = host_ip_input.value
    port = port_input.value
    kai.klonet_login(project_name, user_name, host_ip, port)
    if kai.test_klonet_connection():
        pn.state.notifications.success("Connected to Klonet backend.")
        login_klonet_button.button_type = "default"
        login_klonet_button.disabled = True
    else:
        pn.state.notifications.error('Connect to Klonet backend failed.', duration=0)


# Login to Klonet backend.
login_klonet_button = pn.widgets.Button(name="Login", button_type="primary", width=140)
login_klonet_button.on_click(click_login_klonet_button)


def click_test_klonet_connection_button(event):
    if kai.test_klonet_connection():
        pn.state.notifications.success("The connection has been established.")
    else:
        pn.state.notifications.error("Cannot reach Klonet backend. The reason may be "
                                     "that you are not logged in or the network is unreachable.", duration=0)


# Test Klonet connection.
test_klonet_connection_button = pn.widgets.Button(name="Test Connection", button_type="default", width=140)
test_klonet_connection_button.on_click(click_test_klonet_connection_button)

klonet_card = pn.Card(
    project_name_input,
    user_name_input,
    host_ip_input,
    port_input,
    pn.Row(login_klonet_button, test_klonet_connection_button),
    title="Login to Remote Klonet",
    collapsible=True,
    collapsed=False,
    width=320
)

organizations = ["OpenAI", "Huggingface (Free Trials)", u"阿里通义", u"智谱清言", u"百度文心", u"讯飞星火"]


def update_model_options(event):
    selected_org = event.new
    if selected_org == "OpenAI":
        select_model_button.options = {
            "GPT-3.5 Turbo (4K)": "gpt-3.5-turbo",
            "GPT-3.5 Turbo (16K)": "gpt-3.5-turbo-16k",
            "GPT-4 (8K)": "gpt-4",
            "GPT-4 (32K)": "gpt-4-32k"
        }
    elif selected_org == "Huggingface (Free Trials)":
        select_model_button.options = {
            "StarCoder (Recommended)": "starcoder",
            "CodeLlama 13b (4K)": "codellama-13b-hf",
            "CodeLlama 34b (4K)": "codeLlama-34b-instruct-hf",
            "CodeLlama 34b Phind v2": "phind-codellama-34b-v2"
        }
    elif selected_org == u"阿里通义":
        select_model_button.options = {
            u"通义千问 Turbo": "qwen-turbo",
            u"通义千问 Plus": "qwen-plus"
        }
    elif selected_org == u"智谱清言":
        select_model_button.options = {
            "ChatGLM Lite (8K)": "chatglm_lite",
            "ChatGLM Lite (32K)": "chatglm_lite_32k",
            "ChatGLM Std": "chatglm_std",
            "ChatGLM Pro": "chatglm_pro"
        }
    elif selected_org == u"百度文心":
        select_model_button.options = {
            u"文心一言": "ernie-bot",
            u"文心一言 Turbo": "ernie-bot-turbo"
        }
    elif selected_org == u"讯飞星火":
        select_model_button.options = {
            u"星火认知": "spark-v2"
        }
    else:
        raise NotImplementedError()


# Choose Agent.
select_organization_button = pn.widgets.Select(name="Model Source:", options=[])
select_model_button = pn.widgets.Select(name="Model Name:")
select_organization_button.param.watch(update_model_options, "value")
select_organization_button.options = organizations

APIKey_input = pn.widgets.PasswordInput(
    name="API Key:", placeholder="Default to use the APIKey in key.py", value="")
APPID_input = pn.widgets.PasswordInput(
    name="APP ID (optional):", placeholder="Default to use the APPID in key.py", value="")
APPSecret_input = pn.widgets.PasswordInput(
    name="APP Secret (optional):", placeholder="Default to use the APPSecret in key.py", value="")
run_mode_switch = pn.widgets.Switch(name="Chat-only", value=False)


def click_login_agent_button(event):
    selected_org = select_organization_button.value
    selected_model = select_model_button.value
    api_key = APIKey_input.value
    app_id = APPID_input.value
    app_secret = APPSecret_input.value
    if run_mode_switch.value:
        prompt_template = "The AI assistant tries to be helpful, polite, honest, and humble-but-knowledgeable."
    else:
        prompt_template = "LIKirin/klonetai-prompts"

    try:
        if selected_org == "OpenAI":
            kai.create_agent(
                agent_name="openai", key=api_key, tools=tools, openai_model=selected_model,
                chat_prompt_template=prompt_template, run_prompt_template=prompt_template)
        elif selected_org == "Huggingface (Free Trials)":
            kai.create_agent(
                agent_name=selected_model, key=api_key, tools=tools,
                chat_prompt_template=prompt_template, run_prompt_template=prompt_template)
        elif selected_org == u"阿里通义":
            from agent.dashscope import DashScopeAgent
            custom_agent = DashScopeAgent(
                selected_model, api_key, chat_prompt_template=prompt_template,
                run_prompt_template=prompt_template, additional_tools=tools)
            kai.create_agent(agent=custom_agent)
        elif selected_org == u"智谱清言":
            from agent.chatglm import ChatGLMAgent
            custom_agent = ChatGLMAgent(
                selected_model, api_key, chat_prompt_template=prompt_template,
                run_prompt_template=prompt_template, additional_tools=tools)
            kai.create_agent(agent=custom_agent)
        elif selected_org == u"百度文心":
            from agent.erniebot import ErnieBotAgent
            custom_agent = ErnieBotAgent(
                selected_model, api_key, chat_prompt_template=prompt_template,
                run_prompt_template=prompt_template, additional_tools=tools)
            kai.create_agent(agent=custom_agent)
        elif selected_org == u"讯飞星火":
            from agent.sparkdesk import SparkDesk
            custom_agent = SparkDesk(
                api_key, app_id, app_secret, chat_prompt_template=prompt_template,
                run_prompt_template=prompt_template, additional_tools=tools)
            kai.create_agent(agent=custom_agent)
        else:
            raise NotImplementedError()
    except requests.exceptions.ConnectionError:
        pn.state.notifications.error(f"Connect to {selected_org} ({selected_model}) failed.", duration=0)
        chat_box.append({
            AGENT_NAME: f"Failed to establish connection, try to solve this by using a VPN service."})
        return

    if run_mode_switch.value:
        kai.set_mode("chat")
        chat_box.append(
            {AGENT_NAME: "Chat-only mode is turned on, this chat can only be used for general conversations."})

    login_agent_button.button_type = "success"
    login_agent_button.name = "Try Another Agent"
    login_agent_from_repo_button.name = "Try Another Agent"
    chat_box.append({
        AGENT_NAME: f"I will use the model {selected_model} provided by {selected_org} to take actions."})
    pn.state.notifications.success(f"Connect to {selected_org} ({selected_model}) success.")
    global chat_time
    chat_time = datetime.now().strftime("%Y%m%d-%H%M%S")


login_agent_button = pn.widgets.Button(name="Login", button_type="primary", width=300)
login_agent_button.on_click(click_login_agent_button)

model_card = pn.Card(
    select_organization_button,
    select_model_button,
    APIKey_input,
    APPID_input,
    APPSecret_input,
    login_agent_button,
    title="Login to Remote Agent",
    collapsible=True,
    collapsed=False,
    width=320
)

repo_input = pn.widgets.TextInput(name="HuggingFace Repo:", placeholder="To be done", value="")
login_agent_from_repo_button = pn.widgets.Button(name="Login", button_type="default", width=300)

local_agent_card = pn.Card(
    pn.Column(repo_input, login_agent_from_repo_button),
    title="Login to Local Agent",
    collapsible=True,
    collapsed=False,
    width=320
)

run_mode_card = pn.Card(
    pn.Row(
        pn.widgets.StaticText(value="Run the agent in chat-only mode?"),
        run_mode_switch,
        pn.widgets.TooltipIcon(value='Click "Login" to Remote/Local Agent or "Try Another Agent" to take effect.')
    ),
    title="Agent Mode",
    collapsible=True,
    collapsed=False,
    width=320
)

# Create a KlonetAI agent.
# kai.create_agent(agent_name="starcoder", tools=tools)

USER_NAME = "User"
AGENT_NAME = "KAI"


def command_handler(command):
    if not kai.is_agent_initialized:
        chat_box.append({AGENT_NAME: "Agent not found, please initialize the agent first."})
        return

    global chat_time
    if command == "/reset_chat":
        kai.new_chat()
        chat_box.clear()
        chat_time = datetime.now().strftime("%Y%m%d-%H%M%S")
        chat_box.append({AGENT_NAME: "Welcome to a new chat, let's start over."})
    elif command == "/clear":
        chat_box.clear()
    elif command == "/reset_topo":
        kai.reset_project()
        chat_box.append({AGENT_NAME: "The topology has been reset."})
    elif command == "/chat":
        kai.set_mode("chat")
        chat_box.append({AGENT_NAME: "I'm running in chat mode right now, what do you want to know?"})
    elif command == "/agent":
        kai.set_mode("agent")
        chat_box.append({AGENT_NAME: "I'm running in agent mode now, please tell me what you want to do on Klonet."})
    elif command == "/save":
        save_path = os.path.join(save_dir, f"chat-{chat_time}.txt")
        chat_history = "\n\n".join([f"{row.name}: {row.value[0]}" for row in chat_box.rows])
        fp = open(save_path, "w")
        fp.write(chat_history)
        fp.close()
        chat_box.append({AGENT_NAME: f"This chat has been saved at {save_path}."})
    else:
        chat_box.append({AGENT_NAME: "Invalid command."})


def chat(event):
    if len(chat_box.rows) == 0: return
    last_message = chat_box.rows[-1]
    last_sender = last_message.name
    if last_sender == USER_NAME:
        prompt = last_message.value[0].strip()
        # Handle commands.
        if prompt[0] == '/':
            command_handler(prompt)
            return
        # Check login status.
        if not kai.is_logged_in:
            chat_box.append({AGENT_NAME: "You are not logged in, please login to Klonet first."})
            return
        # Chat with agent.
        with redirect_stdout(io.StringIO()) as output:
            kai.chat(prompt)
        output_str = output.getvalue()
        # Wrap python code.
        start_idx = output_str.find("==Code generated by the agent==")
        if start_idx != -1:
            start_idx += len("==Code generated by the agent==")
            output_str = output_str[:start_idx + 1] + "\n```py\n" + output_str[start_idx + 1:]
        end_idx = output_str.find("==Result==")
        if end_idx != -1:
            for idx in range(end_idx - 1, 0, -1):
                if output_str[idx] != '\n':
                    end_idx = idx
                    break
            output_str = output_str[:end_idx + 1] + "\n```" + output_str[end_idx + 1:]
        chat_box.append({AGENT_NAME: output_str})


chat_box = pn.widgets.ChatBox(
    value=[
        {USER_NAME: "You have joined the chat, please enjoy it!"},
        {AGENT_NAME: "Hello, this is KAI, I can help you interact with Klonet. "
                     "Please tell me what can I do for you :)"}
    ],
    show_names=False
)
chat_box.param.watch(chat, "value")


def click_send_speech_text_button(event):
    speech_text = speech_to_text.value.strip()
    if not speech_text: return
    chat_box.append({USER_NAME: speech_text})


send_speech_text_button = pn.widgets.Button(name="Send", button_type="primary", width=250)
send_speech_text_button.on_click(click_send_speech_text_button)

speech_to_text = SpeechToText(button_type="light")
speech_layout = pn.Column(
    pn.Row(speech_to_text.controls(["value"], jslink=False)),
    pn.Row(send_speech_text_button, speech_to_text)
)
speech_card = pn.Card(
    speech_layout,
    title="Speech to Text",
    collapsible=True,
    collapsed=False,
    width=340
)

quick_command_selector = pn.widgets.Select(
    options=[
        "/agent",
        "/clear",
        "/chat",
        "/reset_chat",
        "/reset_topo",
        "/save",
    ],
    value="/clear",
    width=247
)


def click_send_command_button(event):
    command = quick_command_selector.value
    if not command: return
    chat_box.append({USER_NAME: command})


send_command_button = pn.widgets.Button(name="Send", button_type="primary")
send_command_button.on_click(click_send_command_button)

quick_command_card = pn.Card(
    pn.Row(quick_command_selector, send_command_button),
    title="Quick Command",
    collapsible=True,
    collapsed=False,
    width=340
)


def collect_target_nodes(checkbox, node_input):
    if not kai.is_topo_deployed:
        pn.state.notifications.error(f"No project found, please deploy the project `klonetai` first.")
        chat_box.append({AGENT_NAME: "No project found, please deploy the project `klonetai` first."})
        return []

    all_nodes = kai.remote_topo
    controllers = list(all_nodes["controllers"].keys())
    hosts = list(all_nodes["hosts"].keys())
    routers = list(all_nodes["routers"].keys())
    switches = list(all_nodes["switches"].keys())
    node_groups = checkbox.value
    if "all" in node_groups:
        return controllers + hosts + routers + switches
    else:
        target_nodes = []
        if "hosts" in node_groups: target_nodes += hosts
        if "routers" in node_groups: target_nodes += routers
        if "switches" in node_groups: target_nodes += switches
        if "controllers" in node_groups: target_nodes += controllers

        for item in node_input.value.split(','):
            item = item.strip()
            if item and item not in target_nodes:
                target_nodes.append(item)
        return target_nodes


def click_send_file_button(event):
    if not kai.is_logged_in:
        chat_box.append({AGENT_NAME: "You are not logged in, please login to Klonet first."})
        return

    file_data_list = file_selector.value
    file_name_list = file_selector.filename
    if not (file_data_list and file_name_list): return

    target_nodes = collect_target_nodes(node_group_checkbox, target_node_input)
    if not target_nodes: return

    target_file_path = target_path_input.value
    if not target_file_path: return

    for file_bytes, file_name in zip(file_data_list, file_name_list):
        file_name = file_name.replace(' ', '-')
        for tgt_node in target_nodes:
            if target_file_path[-1] == '/':
                target_file_path = os.path.join(target_file_path, file_name)
            kai.upload_file(tgt_node, (file_name, file_bytes), target_file_path)
    pn.state.notifications.info(f"File uploaded.")
    chat_box.append({AGENT_NAME: f"I have tried to upload the files `{', '.join(file_name_list)}` "
                                 f"to nodes {', '.join(target_nodes)}, but I am not sure whether "
                                 f"they are successfully uploaded. Please manually check them."})
    file_selector.value = None
    file_selector.filename = ""


# File selector
send_file_button = pn.widgets.Button(name="Send", button_type="primary")
send_file_button.on_click(click_send_file_button)

file_selector = pn.widgets.FileInput(multiple=True, width=247)
target_path_input = pn.widgets.TextInput(
    name="Which directory to put these files in?",
    placeholder="/home", value="/home", width=320)
node_group_checkbox_text = pn.widgets.StaticText(
    value="Which types of nodes to upload these files to?")
node_group_checkbox = pn.widgets.CheckBoxGroup(
    inline=True, options=["all", "hosts", "controllers", "routers", "switches"])
target_node_input = pn.widgets.TextInput(
    name="Which nodes to upload these files to?", placeholder="h1,h2,h4", value="")

file_selector_card = pn.Card(
    pn.Column(
        target_path_input,
        node_group_checkbox_text,
        node_group_checkbox,
        target_node_input,
        pn.Row(file_selector, send_file_button),
    ),
    title="Upload File",
    collapsible=True,
    collapsed=False,
    width=340
)

SOURCE_COMMAND = "source /root/.bashrc"
preset_command = f"{SOURCE_COMMAND}"


def click_run_command_button(event):
    if not kai.is_logged_in:
        chat_box.append({AGENT_NAME: "You are not logged in, please login to Klonet first."})
        return

    command = command_input.value.strip()
    if not command: return

    target_nodes = collect_target_nodes(node_group_checkbox_for_command, target_node_input_for_command)
    if not target_nodes: return

    chat_box.append({AGENT_NAME: f"Run commands `{command}` on nodes {', '.join(target_nodes)}."})

    global preset_command
    ctns = {"list_type": "specified_ctn_list", "list": target_nodes}

    full_command = preset_command[:]
    conda_path = conda_path_input.value
    conda_env = conda_env_name.value
    if conda_path and conda_env:
        full_command = " && ".join([
            full_command,
            f"source {conda_path}",
            f"conda activate {conda_env}"
        ])
    full_command = " && ".join([full_command, command])
    print(f"Running command: {full_command}")
    response = kai.batch_exec(ctns, f'bash -c "{full_command}"', block="true")

    # Update current path.
    if ';' in command:
        command_list = command.split(';')
    elif '&&' in command:
        command_list = command.split('&&')
    else:
        command_list = [command]
    for cmd in command_list:
        cmd = cmd.strip()
        if cmd[:4] == "cd /" or cmd == "cd":
            preset_command = f"{SOURCE_COMMAND} && {cmd}"
        elif cmd[:3] == "cd ":
            if preset_command == SOURCE_COMMAND:
                preset_command = f"{SOURCE_COMMAND} && {cmd}"
            else:
                preset_command = " && ".join([preset_command, cmd])

    def reformat_json(response):
        reformated_response = {}
        for worker_ip, item in response.items():
            result = item["worker_exec_results"]
            for node_name, node_item in result.items():
                reformated_response[node_name] = node_item["output"]
        return reformated_response

    if type(response) is not dict:
        chat_box.append({AGENT_NAME: response})
        return

    output_str = ""
    for node_name, output in reformat_json(response).items():
        output_str += f"```shell\n"
        output_str += f"{node_name}$ {command}"
        output_str += f"\n{output}"
        output_str += "\n```\n"
    chat_box.append({AGENT_NAME: output_str})


def click_clear_command_button(event):
    command_input.value = ""


command_input = pn.widgets.TextAreaInput(name="Command:", value="", width=320, height=70)
run_command_button = pn.widgets.Button(name="Run", button_type="primary", width=150)
run_command_button.on_click(click_run_command_button)
clear_command_button = pn.widgets.Button(name="Clear Command", button="default", width=150)
clear_command_button.on_click(click_clear_command_button)
node_group_checkbox_text_for_command = pn.widgets.StaticText(
    value="Which types of nodes to run these commands?")
node_group_checkbox_for_command = pn.widgets.CheckBoxGroup(
    inline=True, options=["all", "hosts", "controllers", "routers", "switches"])
target_node_input_for_command = pn.widgets.TextInput(
    name="Which nodes to run these commands?", placeholder="h1,h2,h4", value="", width=320)


def configure_conda(event):
    kai.additional_info["conda_path"] = conda_path_input.value
    kai.additional_info["conda_env"] = conda_env_name.value


conda_path_input = pn.widgets.TextInput(
    name="Path to conda configuration file within containers:",
    placeholder="/opt/conda/etc/profile.d/conda.sh",
    value="/opt/conda/etc/profile.d/conda.sh")
conda_path_input.param.watch(configure_conda, "value")

conda_env_name = pn.widgets.TextInput(
    name="Name of conda environment to be used:", placeholder="base", value="base")
conda_path_input.param.watch(configure_conda, "value")
configure_conda(None)

# Execute batch commands
exec_command_card = pn.Card(
    pn.Column(
        command_input,
        node_group_checkbox_text_for_command,
        node_group_checkbox_for_command,
        target_node_input_for_command,
        conda_path_input,
        conda_env_name,
        pn.Row(run_command_button, clear_command_button)
    ),
    title="Command Execution",
    collapsible=True,
    collapsed=False,
    width=340
)

layout_left = pn.Column(klonet_card, model_card, local_agent_card, run_mode_card)
layout_right = pn.Column(speech_card, exec_command_card, file_selector_card, quick_command_card)
layout = pn.Row(layout_left, chat_box, layout_right)

pn.serve(
    layout,
    title="KAI Dashboard",
    websocket_max_message_size=1024 * 1024 * 1014,
    http_server_kwargs={"max_buffer_size": 1024 * 1024 * 1014}
)

chatbox.py
Lizonghang/KlonetAI
# Not all the API Keys below are required, just choose one of them.

# Apply on OpenAI: https://platform.openai.com/docs/guides/production-best-practices/api-keys
OpenAI_API_Key = "<Your-OpenAI-API-Key>"

# Apply on ZhuPuAI: https://open.bigmodel.cn/dev/api#auth
ZhiPuAI_API_Key = "<Your-ZhiPuAI-API-Key>"

# Apply on Hugging Face: https://huggingface.co/docs/hub/spaces-sdks-docker-argilla#set-up-passwords-and-api-keys-using-secrets-optional
Huggingface_API_Key = "<Your-Huggingface-API-Key>"

# Apply on Baidu AI Studio: https://aistudio.baidu.com/usercenter/token
AI_Studio_APT_Key = "<Your-AIStudio-API-Key>"

# Apply on Aliyun DashScope: https://help.aliyun.com/zh/dashscope/developer-reference/activate-dashscope-and-create-an-api-key?spm=a2c4g.11186623.0.0.90a174faPZtci1
DashScope_API_Key = "<Your-DashScope-API-Key>"

# Apply on Xunfeiyun: https://www.xfyun.cn/doc/spark/general_url_authentication.html
SparkDesk_AppId = "<Your-SparkDesk-APPID>"
SparkDesk_API_Secret = "<Your-SparkDesk-APISecret>"
SparkDesk_API_Key = "<Your-SparkDesk-APIKey>"

key.py
Lizonghang/KlonetAI
"""
This is vemu api.
"""
from .image import ImageManager
from .link import LinkManager
from .node import NodeManager
from .project import ProjectManager
from .cmd import CmdManager
from .common.base_classes import Node, Image, Link, Topo, LinkConfiguration
from .common.errors import *
klonet_api/__init__.py
Lizonghang/KlonetAI
from .common import Manager

class CmdManager(Manager):
    '''命令管理类

    用于在已创建项目的节点中执行shell命令。

    Attributes:
        user(str): 用户名
        project(str): 项目名
    '''
    def __init__(self, user_name, project_name,
                backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name
        self.project = project_name

    def exec_cmds_in_nodes(self, node2cmds, block="false", timeout=60):
        """在多个节点中执行多条shell命令

        需注意，若执行较长时间阻塞的命令，如iperf3 -s，后台会执行此命令，
        但会触发超时机制，无法获取到命令的退出码及输出。

        Args:
            node2cmd(dict): 命令字典，key为节点名，value为shell命令列表。如：
                {"h1": ["ls", "ls"],"h2": ["pwd"]}

        Returns:
            一个字典，描述了节点命令执行的情况。key为节点名，value为该节点中
            具体命令的执行结果字典。比如：

            {'h1': {'ls': { "exit_code": 0, "output": "bin"}}}

        """
        payload = {
            "user": self.user,
            "topo": self.project,
            "node_and_cmd": node2cmds,
            "block": block,
            "cmd_timeout_s": timeout
        }
        resp = self._post("/master/node_exec_cmd/", json=payload)
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        return resp_json["exec_results"]




klonet_api/cmd.py
Lizonghang/KlonetAI
from .base_funcs import *
from .base_classes import *
from .errors import *
klonet_api/common/__init__.py
Lizonghang/KlonetAI
import requests
import copy
from .. import config
from .base_funcs import cidr2ip_and_netmask, get_plural_of_words
from .errors import *


'''基础类'''

class Manager(object):
    '''vemu api的manager基类
    
    提供：1. 后端ip和端口的配置。
    2. 基本的post/get等请求方式的封装。
    3. 对response的解析。

    Attributes:
        backend_ip(str): 后端服务器IP
        backend_port(int): 后端服务器端口
        url(str): 请求url
    '''
    def __init__(self, backend_ip=None, backend_port=None):
        if not backend_ip or not backend_port:
            try:
                backend_ip, backend_port = config.backend_ip, config.backend_port
            except AttributeError:
                raise ValueError("Please config backend_ip and backend_port by "
                "function args or config.py!")
        self.url = (f"http://{backend_ip}:{backend_port}")

    def _post(self, url_suffix, json=None, data=None):
        return requests.post(url=f"{self.url}{url_suffix}", 
            json=json, data=data)

    def _delete(self, url_suffix, json=None, data=None):
        return requests.delete(url=f"{self.url}{url_suffix}", 
            json=json, data=data)

    def _put(self, url_suffix, json=None, data=None):
        return requests.put(url=f"{self.url}{url_suffix}", 
            json=json, data=data)

    def _get(self, url_suffix, json=None, data=None, params=None):
        return requests.get(url=f"{self.url}{url_suffix}", 
            json=json, data=data, params=params)

    def _parse_resp(self, response):
        '''对response对象进行解析，返回其json格式。

        同时该函数会对response的状态码进行校验，若状态码不是200，则会抛出异常。

        Args:
            response(Response): requests库的Response对象

        Returns:
            返回response.json()

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
        '''
        if response.status_code != 200:
            raise HttpStatusError(f"HTTP status code is not 200. "
                f"response.status_code = {response.status_code}, "
                f"response.text = {response.text}")

        response_json = None
        try:
            response_json = response.json()
        except requests.exceptions.JSONDecodeError:
            raise JsonDecodeError(f"The response does not contain valid json. "
                f"response.status_code = {response.status_code}, "
                f"response.text = {response.text}")

        return response_json

    def _check_resp_code(self, resp_json, code_field="code", success_code=1,
                        msg_filed="msg"):
        '''检查response json中的code

        若code不为success_code，则抛出异常

        Args:
            resp_json(dict): 响应response的json，例子:
                {"code": 1,
                "msg": "success!"
                }
            code_filed(str): 响应response的json中记录code的key
            success_code(int): 执行成功的code
            msg_filed(str): 响应response的json中记录提示信息的key

        Raises:
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
        '''
        if resp_json[code_field] != success_code:
            raise VemuExecError(f"Return code={resp_json[code_field]} after vemu "
                f"execute this request, msg: {resp_json[msg_filed]}")

class Dict2Class(object):
    '''镜像、节点、链路等的基类'''
    # https://stackoverflow.com/a/1305663
    def __init__(self, **properties):
        self.__dict__.update(properties)

    def dictform(self):
        return self.__dict__

class Image(Dict2Class):
    '''镜像类

    需要传字典来确定其拥有的属性，要查看实例化对象所具有的属性，请调用dictform()方法
    '''
    def __init__(self, **properties):
        super().__init__(**properties)

class Node(Dict2Class):
    '''节点类

    需要传字典来确定其拥有的属性要查看实例化对象所具有的属性，请调用dictform()方法
    '''
    def __init__(self, **properties):
        super().__init__(**properties)

class Link(Dict2Class):
    '''链路类

    需要传字典来确定其拥有的属性要查看实例化对象所具有的属性，请调用dictform()方法
    '''
    def __init__(self, **properties):
        super().__init__(**properties)

class Topo(Dict2Class):
    '''拓扑类

    用于项目创建前通过类中的add_node和add_link方法设计拓扑，设计完成后需将Topo对
    象传入TopoManager的deploy方法，以完成实际的项目创建。
    '''
    def __init__(self, **properties):
        self.__dict__ = {} # topo结构的维护是由拓扑对象中的字典来实现
        elements = ["links", "controllers", "hosts", "routers", "switches"]
        for element in elements:
            self.__dict__.setdefault(element, {})
        super().__init__(**properties)

    def add_node(self, image, node_name=None, resource_limit=None,
            location={"x": 0, "y": 0}, worker_specified=None):
        '''向Topo对象中添加节点。

        Args:
            image(Image): 节点所用镜像的Image对象，需先通过镜像相关API获取
            node_name(str): 所添加节点的名称。
            resource_limit(dict): 资源限制。例子：
                {"cpu": "1000", # CPU利用率限制，单位：%
                "mem": "1000" # 内存限制，单位：Mbytes
                }
                默认为None，将采用镜像中的默认资源限制。
            location(dict): 在前端画布上的横纵坐标，x和y均应该大于等于0。例子：
                {"x": 0,
                "y": 0
                }
                默认值同样为{"x":0, "y":0}

        Returns:
            所添加节点的Node对象，便于后续直接对其进行相关操作

        Raises:
            NodeDuplicatesError: 当所添加的节点名重复时，触发此异常
        '''   
        # 默认名称分配
        if not node_name:
            node_name = self._assign_default_node_name()

        # 名称重复检查
        nodes = self.get_nodes()
        if node_name in nodes.keys():
            raise NodeDuplicatesError(f"node name [{node_name}] is duplicate, "
                f"existing node names are {list(nodes.keys())}")

        # node对象构建
        node = copy.deepcopy(image) # 将image对象当作node对象使用
        if resource_limit:
            node.resource_limit = resource_limit
        node.name = node_name
        node.x = location["x"]
        node.y = location["y"]

        if worker_specified:
            node.config.update({"worker_specified": worker_specified})

        # 加入topo的字典中
        category = get_plural_of_words(node.type)
        self.__dict__[category][node_name] = node.dictform()

        return Node(**node.dictform())

    def add_link(self, src_node, dst_node, link_name=None, src_IP="",
        dst_IP=""):
        '''向Topo对象中添加链路

        参数中的源和目的仅用于区分不同的两端，不含有方向的意思，链路为双向链路

        Args:
            src_node(Node): 源节点的Node对象
            dst_node(Node): 目的节点的Node对象
            link_name(str): 要添加的链路名，例如"link1"，默认为None，默认情况下会将链路命
                名为"l<拓扑中现有链路数量+1>"
            src_IP(str): 源节点的IP地址，例如"192.168.1.1/24"，默认为""
            dst_IP(str): 目的节点的IP地址，例如"192.168.1.2/24"，默认为""

        Returns:
            所添加链路的Link对象，便于后续直接对其进行相关操作
        '''
        # 参数检查
        if src_node.name == dst_node.name:
            raise ValueError(f"Node cannot connect to itself!")
        if src_IP != "":
            cidr2ip_and_netmask(src_IP)
        if dst_IP != "":
            cidr2ip_and_netmask(dst_IP)

        # 默认名称分配
        if not link_name:
            link_name = self._assign_default_link_name()

        # 检查平行边
        self._check_parallel_link(link_name, src_node.name, dst_node.name)

        # 名称重复检查
        links = self.get_links()
        if link_name in links.keys():
            raise LinkDuplicatesError(f"link name [{link_name}] is duplicate, "
                f"existing link names are {list(links.keys())}")

        # 构建Link对象并将其加入Topo中
        link = Link()
        link.config = {
            "source" : {
                "bw_kbit" : "",
                "queue_size_byte" : "",
                "delay_us" : "",
                "loss_rate" : "", 
                "jitter_us" : "", 
                "correlation" : "", 
                "delay_distribution" : "normal"
            },
            "target":{
                "bw_kbit" : "",
                "queue_size_byte" : "",
                "delay_us" : "",
                "loss_rate" : "", 
                "jitter_us" : "", 
                "correlation" : "", 
                "delay_distribution" : "normal"
            }
        }
        link.name = link_name
        link.source = src_node.name
        link.sourceIP = src_IP
        link.sourceType = src_node.type
        link.target = dst_node.name
        link.targetIP = dst_IP
        link.targetType = dst_node.type

        self.__dict__["links"][link.name] = link.dictform()

        # 修改节点信息
        if src_IP != "":
            ip, netmask = cidr2ip_and_netmask(src_IP)
            nic_nickname = f"{src_node.name}{dst_node.name}"
            category = get_plural_of_words(src_node.type)
            self.__dict__[category][src_node.name]["interfaces"].append(
                {"ip": ip, "netmask": netmask, "name": nic_nickname})

        if dst_IP != "":
            ip, netmask = cidr2ip_and_netmask(dst_IP)
            nic_nickname = f"{dst_node.name}{src_node.name}"
            category = get_plural_of_words(dst_node.type)
            self.__dict__[category][dst_node.name]["interfaces"].append(
                {"ip": ip, "netmask": netmask, "name": nic_nickname})

        return link

    def get_nodes(self):
        '''获取Topo对象中所有的节点名及对应的Node对象

        Returns:
            一个字典，包含该Topo对象中所有的节点名及对应的Node对象。例子：
            {"h1": h1的Node对象,
            ...
            }
        '''
        nodes = {}
        for type in self.__dict__.keys():
            if type == "links":
                continue
            for node_name, node_dict in self.__dict__[type].items():
                nodes[node_name] = Node(**node_dict)
        
        return nodes

    def get_links(self):
        '''
        获取Topo对象中所有的链路名及对应的Link对象

        Returns:
            一个字典，包含该Topo对象中所有的链路名及对应的Link对象。例子：

            {"l1": l1的Node对象,
            ...
            }
        '''
        links = {}
        for link_name, link_dict in self.__dict__["links"].items():
            links[link_name] = Link(**link_dict)
        
        return links

    def _assign_default_link_name(self):
        '''分配默认链路名

        默认链路名为"l<拓扑中现有链路数量+1>"
        '''
        link_num = len(self.get_links())
        link_name = f"l{link_num+1}"
        return link_name

    def _assign_default_node_name(self):
        '''分配默认节点名

        默认节点名为"n<拓扑中现有链路数量+1>"
        '''
        node_num = len(self.get_nodes())
        node_name = f"n{node_num+1}"
        return node_name

    def _check_parallel_link(self, link_name, src_node_name, dst_node_name):
        '''检查新添加的链路是否为平行边（即新链路与已有链路的两端节点名相同）

        Args:
            src_node_name(str): 链路的源节点名
            dst_node_name(str): 链路的目的节点名

        Raises:
            LinkParallelError: 当出现平行边（即新边与已有边的两端节点名相同）时，触发
                该异常
        '''
        links = self.get_links()
        my_endpoints = set([src_node_name, dst_node_name])
        for _, link in links.items():
            exist_endpoints = set([link.source, link.target])
            if my_endpoints == exist_endpoints:
                raise LinkParallelError(f"New link {link_name}({src_node_name}"
                    f"---{dst_node_name}) repeat with exist link {link.name}"
                    f"({link.source}---{link.target})")

class LinkConfiguration(Dict2Class):
    '''链路配置类。

    默认的属性及说明为：
    "bw_kbps":"10000", # 链路带宽（kbps）,需为正数
    "delay_us":"0", # 链路时延(us)，需为非负数
    "jitter_us":"0", # 时延抖动(us)，需为非负数
    "correlation":"0%", # 抖动相关率(%)，需为非负百分率，如1%
    "delay_distribution":"uniform", # 时延抖动分布，可选项：uniform/normal/pareto/paretonomal
    "loss":"0", # 链路丢包率(%)，需为非负数，如1
    "queue_size_bytes":"100000", # 队列大小(字节)，需为非负数，如10000
    "linkchoice":"static", # 固定为static，无需改变
    "link":None, # 链路名
    "ne":None # 节点名
    '''
    def __init__(self, **properties):
        default_config = {
            "bw_kbps":"10000", # 链路带宽（kbps）,需为正数
            "delay_us":"0", # 链路时延(us)，需为非负数
            "jitter_us":"0", # 时延抖动(us)，需为非负数
            "correlation":"0%", # 抖动相关率(%)，需为非负百分率，如1%
            "delay_distribution":"uniform", # 时延抖动分布，可选项：uniform/normal/pareto/paretonomal
            "loss":"0", # 链路丢包率(%)，需为非负数，如1
            "queue_size_bytes":"100000", # 队列大小(字节)，需为非负数，如10000
            "linkchoice":"static", # 固定为static，无需改变
            "link":None, # 链路名
            "ne":None # 节点名
        }
        self.__dict__ = default_config
        super().__init__(**properties)
klonet_api/common/base_classes.py
Lizonghang/KlonetAI
import re

'''基础函数'''
def is_ip_leagal(ip):
    '''检查ip地址的合法性

    Args:
        ip(str): ip地址，如192.168.1.1（合法），256.0.0.1（不合法），aaaa（不合法）

    Returns:
        合法则返回True，不合法返回False
    '''
    compile_ip = re.compile("^(1\d{2}|2[0-4]\d|25[0-5]|[1-9]\d|[1-9])\."
        "(1\d{2}|2[0-4]\d|25[0-5]|[1-9]\d|\d)\."
        "(1\d{2}|2[0-4]\d|25[0-5]|[1-9]\d|\d)\."
        "(1\d{2}|2[0-4]\d|25[0-5]|[1-9]\d|\d)$")
    if compile_ip.match(ip):
        return True
    else:    
        return False

def is_cidr_leagal(cidr):
    '''检查cidr形式的ip地址和掩码的合法性

    Args:
        ip(str): cidr形式的ip地址，如192.168.1.1/24（合法），256.0.0.1（不合法），
            aaaa（不合法），192.168.1.1/33（不合法）
    '''
    ip_and_mask = cidr.split("/")

    if len(ip_and_mask) != 2:
        return False
    else:
        try:
            mask = int(ip_and_mask[1])
        except:
            return False
        if mask < 0 or mask > 32:
            return False
        if not is_ip_leagal(ip_and_mask[0]):
            return False

    return True

def cidr2ip_and_netmask(cidr):
    '''将cidr形式的地址转换为ip和子网掩码

    如192.168.1.1/24转换为192.168.1.1和255.255.255.0。
    也可作为cidr形式的合法性检验函数。

    Args:
        cidr(str): cidr形式的地址，如192.168.1.1/24

    Returns:
        返回转换后的ip及子网掩码如192.168.1.1,255.255.255.0

    Raises:
        ValueError: cidr形式的地址不合法
    '''
    if not is_cidr_leagal(cidr):
        raise ValueError(f"Address [{cidr}] is illegal, please check!")

    ip_and_netmask = cidr.split("/")
    ip = ip_and_netmask[0]
    netmask = cidr_netmask(int(ip_and_netmask[1]))

    return ip, netmask

def cidr_netmask(prefix):
    '''将int类型的掩码转换为地址类型的掩码

    Args:
        prefix(int): int类型的掩码，如24

    Returns:
        地址类型的掩码，如255.255.255.0
    '''
    bin_arr = ['0' for _ in range(32)]
    for i in range(prefix):
        bin_arr[i] = '1'
    tmpmask = [''.join(bin_arr[i * 8:i * 8 + 8]) for i in range(4)]
    tmpmask = [str(int(tmpstr, 2)) for tmpstr in tmpmask]
    return '.'.join(tmpmask)

def get_plural_of_words(word):
    """返回单词的复数形式

    用于根据节点/镜像的type字段返回复数形式type的字段。

    Args:
        word(str): type字段

    Returns:
        word的复数形式
    """
    if word in ['host', 'router', 'floodlight', 'controller', 'dpdk']:
        return word + 's'
    elif word in ['switch',]:
        return word + 'es'
    else:
        raise TypeError(f'get_plural_of_words() do not support the word '
            f'[{word}]')
klonet_api/common/base_funcs.py
Lizonghang/KlonetAI
'''异常定义'''

class HttpStatusError(RuntimeError):
    '''当HTTP的返回状态码不为200时，触发此异常'''
    pass

class JsonDecodeError(RuntimeError):
    '''当返回体不包含json时，触发此异常'''
    pass

class VemuExecError(RuntimeError):
    '''当HTTP请求成功，但json中的返回码不为1时，触发此异常'''
    pass

class NodeNotExistsError(RuntimeError):
    pass

class NodeDuplicatesError(RuntimeError):
    pass

class LinkDuplicatesError(RuntimeError):
    pass

class LinkNotExistsError(RuntimeError):
    '''当目标链路不存在时，触发此异常'''
    pass

class LinkParallelError(RuntimeError):
    '''当出现平行边（即新边与已有边的两端节点名相同）时，触发该异常'''
    pass

class LinkInconsistentError(RuntimeError):
    '''当链路属性配置时，链路两端的LinkConfiguration对象的链路名不一致时，触发该异常
    '''
    pass
klonet_api/common/errors.py
Lizonghang/KlonetAI
"""配置文件

可对后端服务器的IP及端口进行配置

"""
#: str: 后端服务器IP
backend_ip = "220.243.137.82"
#: int: 后端服务器端口
backend_port = 12352

klonet_api/config.py
Lizonghang/KlonetAI
from .common import Manager, Image

class ImageManager(Manager):
    '''镜像管理类

    Attributes:
        user(str): 用户名
    '''
    def __init__(self, user_name, backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name

    def get_images(self, quiet=False):
        '''获取当前用户的所有镜像名及镜像对象。

        Args:
            quiet(bool): 默认为False。若为False，则将打印镜像名列表；否则将关闭打印。

        Returns:
            一个字典，key为镜像名，value为Image对象
        '''
        resp = self._get("/my/image/", params={"username": self.user})
        resp_json = self._parse_resp(resp)
        try:
            self._check_resp_code(resp_json)
        except KeyError:
            pass
        image_list = {}
        for registry_type in resp_json.keys(): # 镜像仓库类型，如private/public
            for type in resp_json[registry_type].keys():# 镜像类型，如host/switch
                for image_dict in resp_json[registry_type][type]:
                    # 注意这里的subtype为前端显示的镜像名，在这里同样用作镜像名来向
                    # 用户屏蔽细节
                    image_list[image_dict["subtype"]] = Image(**image_dict)

        if not quiet:
            print(f"{self.user} has these images: {list(image_list.keys())}")
            
        return image_list
klonet_api/image.py
Lizonghang/KlonetAI
from .common import Manager, Link, LinkNotExistsError, LinkParallelError, LinkInconsistentError, cidr2ip_and_netmask

class LinkManager(Manager):
    '''链路管理类

    用于已创建项目中的链路相关管理。

    Attributes:
        user(str): 用户名
        project(str): 项目名

    '''
    def __init__(self, user_name, project_name,
                backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name
        self.project = project_name

    def dynamic_add_link(self, link_name, src_node, dst_node, src_IP="",
        dst_IP=""):
        '''动态添加链路

        注意：该API仅对已创建项目生效！
        （参数中的源和目的仅用于区分不同的两端，不含有方向的意思，链路为双向链路）

        Args:
            link_name(str): 要添加的链路名
            src_node(Node): 源节点的Node对象
            dst_node(Node): 目的节点的Node对象
            src_IP(str): 源节点的IP地址，例如"192.168.1.1/24"，默认为""
            dst_IP(str): 目的节点的IP地址，例如"192.168.1.2/24"，默认为""

        Returns:
            None

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
            LinkParallelError: 当出现平行边（即新边与已有边的两端节点名相同）时，触发
                此异常
        '''
        # 参数检查
        if src_node.name == dst_node.name:
            raise ValueError(f"Node cannot connect to itself!")
        if src_IP != "":
            cidr2ip_and_netmask(src_IP)
        if dst_IP != "":
            cidr2ip_and_netmask(dst_IP)

        # 检查平行边
        self._check_parallel_link(link_name, src_node.name, dst_node.name)

        # 创建链路
        payload = {
            "user": self.user,
            "topo": self.project,
            "info": {
                "config": {
                    "source": {"bw_kbit":"", "queue_size_byte":"", "delay_us":"",
                        "loss_rate":"", "jitter_us":"", "correlation":"",
                        "delay_distribution":"normal"},
                    "target":{"bw_kbit":"", "queue_size_byte":"", "delay_us":"",
                        "loss_rate":"", "jitter_us":"", "correlation":"",
                        "delay_distribution":"normal"}
                }
            }
        }
        payload["info"]["name"] = link_name
        payload["info"]["source"] = src_node.name
        payload["info"]["sourceIP"] = src_IP
        payload["info"]["sourceType"] = src_node.type
        payload["info"]["target"] = dst_node.name
        payload["info"]["targetIP"] = dst_IP
        payload["info"]["targetType"] = dst_node.type

        resp = self._post("/modification/link/", json=payload)
        self._check_resp_code(self._parse_resp(resp))

        # 修改节点信息
        if src_IP != "":
            ip, netmask = cidr2ip_and_netmask(src_IP)
            print(netmask)
            nic_nickname = f"{src_node.name}{dst_node.name}"
            src_node.interfaces.append({"ip": ip, "netmask": netmask, 
                "name": nic_nickname})
            payload = {}
            payload = {"user": self.user, "topo": self.project, 
                "info": src_node.dictform()}
            resp = self._put("/modification/container/", json=payload)
            self._check_resp_code(self._parse_resp(resp))

        if dst_IP != "":
            ip, netmask = cidr2ip_and_netmask(dst_IP)
            nic_nickname = f"{dst_node.name}{src_node.name}"
            dst_node.interfaces.append({"ip": ip, "netmask": netmask, 
                "name": nic_nickname})
            payload = {}
            payload = {"user": self.user, "topo": self.project, 
                "info": dst_node.dictform()}
            resp = self._put("/modification/container/", json=payload)
            self._check_resp_code(self._parse_resp(resp))

    def dynamic_delete_link(self, link_name):
        '''动态删除节点。
        
        注意：该API仅对已创建项目生效！

        Args:
            node_name(str): 要删除节点名

        Returns:
            None

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
        '''
        link = self.get_link(link_name)

        payload = {"user": self.user, "topo": self.project, 
            "info": link.__dict__}
        resp = self._delete("/modification/link/", json=payload)
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

    def config_link(self, src_link_config, dst_link_config):
        '''配置链路属性。

        链路属性底层使用linux traffic control(tc)实现，因此链路属性实质上是对链路两侧
        节点上网卡的队列的配置，从网卡发出的数据包将经过配置的队列。

        Args:
            src_link_config(LinkConfiguration): 链路源端的LinkConfiguration对象
            dst_link_config(LinkConfiguration): 链路目的端的LinkConfiguration对象

        Returns:
            None

        Raises:
            LinkInconsistentError: 当链路属性配置时，链路两端的LinkConfiguration
                对象的链路名不一致时，触发该异常。
        '''
        # 是否是同一条链路检查
        if src_link_config.link != dst_link_config.link:
            raise LinkInconsistentError(f"src_link_config.link("
                f"{src_link_config.link}) is inconsistent with "
                f"dst_link_config.link({dst_link_config.link}), please check!")

        src_link_config.link = f"link_{src_link_config.link}"
        dst_link_config.link = f"link_{dst_link_config.link}"
        payload = {
            "user": self.user,
            "topo": self.project,
            "links": [src_link_config.dictform(), dst_link_config.dictform()]
        }

        resp = self._post('/master/link/', json=payload)
        print(self._parse_resp(resp))
        self._check_resp_code(self._parse_resp(resp))
        

    def clear_link_configuration(self, link_name):
        '''清除链路上的队列配置。

        Args:
            link_name(str): 链路名

        Returns:
            None
        '''
        link = self.get_link(link_name)
        payload = {
            "user": self.user,
            "topo": self.project,
            "links": [
                {
                    "link": f"link_{link_name}",
                    "linkchoice": "static",
                    "ne": link.source
                },
                {
                    "link": f"link_{link_name}",
                    "linkchoice": "static",
                    "ne": link.target
                }
            ]
        }
        resp = self._delete("/master/link/", json=payload)
        self._check_resp_code(self._parse_resp(resp))

    def get_links(self):
        '''获取该项目下所有的链路名及对应的Link对象

        Returns:
            一个字典，包含该项目所有的链路名及对应的Link对象。比如：

            {"l1": l1的Link对象,"l2": l2的Link对象}

        '''
        resp = self._get(f"/re/project/{self.project}/",
            params={"user": self.user})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        links = {}
        for link_name, link_dict in resp_json["project"]["topo"][
            "links"].items():
            links[link_name] = Link(**link_dict)

        return links

    def get_link(self, link_name):
        '''获取目标链路的Link对象。

        Args:
            link_name(str): 链路名

        Returns:
            目标链路的Link对象。

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
            LinkNotExistsError: 当目标链路不存在时，触发此异常
        '''
        links = self.get_links()
        try:
            link = links[link_name]
            return link
        except KeyError:
            raise LinkNotExistsError(f"Link [{link_name}] does not exist, "
                f"avaliable links are {list(links.keys())}")

    def _check_parallel_link(self, link_name, src_node_name, dst_node_name):
        '''检查新添加的链路是否为平行边（即新链路与已有链路的两端节点名相同）

        Args:
            src_node_name(str): 链路的源节点名
            dst_node_name(ste)：链路的目的节点名

        Returns:
            None。

        Raises:
            LinkParallelError: 当出现平行边（即新边与已有边的两端节点名相同）时，触发
                该异常
        '''
        links = self.get_links()
        my_endpoints = set([src_node_name, dst_node_name])
        for _, link in links.items():
            exist_endpoints = set([link.source, link.target])
            if my_endpoints == exist_endpoints:
                raise LinkParallelError(f"New link {link_name}({src_node_name}"
                    f"---{dst_node_name}) repeat with exist link {link.name}"
                    f"({link.source}---{link.target})")


klonet_api/link.py
Lizonghang/KlonetAI
from .common import Manager, Node, NodeNotExistsError
import copy

class NodeManager(Manager):
    """节点管理类

    Attributes:
        user(str): 用户名
        project(str): 项目名

    """
    def __init__(self, user_name, project_name,
                backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name
        self.project = project_name

    def dynamic_add_node(self, node_name, image, resource_limit=None,
            location={"x": 0, "y": 0}, worker_specified=None):
        '''动态添加节点。

        注意：该API仅对已创建项目生效！

        Args:
            node_name(str): 所添加节点的名称
            image(Image): 节点所用镜像的Image对象，需先通过镜像相关API获取
            resource_limit(dict): 资源限制。例子：
                {"cpu": "1000", # CPU利用率限制，单位：%
                "mem": "1000" # 内存限制，单位：Mbytes}
                默认为None，将采用镜像中的默认资源限制。
            location(dict): 在前端画布上的横纵坐标，x和y均应该大于等于0。例子：
                {"x": 0,
                "y": 0}
                默认值同样为{"x":0, "y":0}
            worker_specified(str): 指定动态添加节点时的宿主机，如172.17.0.16（平台
                使用宿主机的ip地址来表示宿主机）；若不进行指定，则将由平台自动进行指
                定。

        Returns:
            所添加节点的Node对象，便于后续直接对其进行相关操作

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
        '''
        node = copy.deepcopy(image) # 将image对象当作node对象使用
        if resource_limit:
            node.resource_limit = resource_limit
        node.name = node_name
        node.x = location["x"]
        node.y = location["y"]
        if worker_specified:
            node.config.update({"worker_specified": worker_specified})

        payload = {"user": self.user, "topo": self.project,
            "info": node.dictform()}
        resp = self._post("/modification/container/", json=payload)
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        return Node(**node.dictform())

    def dynamic_delete_node(self, node_name):
        '''动态删除节点。

        注意：该API仅对已创建项目生效！

        Args:
            node_name: 要删除节点名

        Returns:
            None

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
        '''
        node = self.get_node(node_name)

        payload = {"user": self.user, "topo": self.project, 
            "info": node.__dict__}
        resp = self._delete("/modification/container/", json=payload)
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

    # def dynamic_modify_node(self, node):
    #     '''
    #     TODO: 灵活性太多、输入参数不确定？
    #     '''
    #     pass

    def get_nodes(self):
        '''获取该项目下所有的节点名及对应的Node对象

        Returns:
            一个字典，包含该项目所有的节点名及对应的Node对象。例子：
            
            {"h1": h1的Node对象,
            "h2": h2的Node对象,
            "...": ...}

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
        '''
        resp = self._get(f"/re/project/{self.project}/",
            params={"user": self.user})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        nodes = {}
        for type in resp_json["project"]["topo"].keys():
            if type == "links":
                continue
            for node_name, node_dict in resp_json["project"]["topo"][
                type].items():
                nodes[node_name] = Node(**node_dict)

        return nodes

    def get_node(self, node_name):
        '''获取目标节点的Node对象。

        Args:
            node_name(str): 节点名

        Returns:
            目标节点的Node对象。

        Raises:
            HttpStatusError: 当HTTP的返回状态码不为200时，触发此异常
            JsonDecodeError: 当返回体不包含json时，触发此异常
            VemuExecError: 当HTTP请求成功，但json中的返回码不为1时，触发此异常
            NodeNotExistsError: 当目标节点不存在时，触发此异常
        '''
        nodes = self.get_nodes()
        try:
            node = nodes[node_name]
            return node
        except KeyError:
            raise NodeNotExistsError(f"Node [{node_name}] does not exist, "
                f"avaliable nodes are {list(nodes.keys())}")


    def ssh_service(self, node_name, start, passwd="123456"):
        """对该项目下一个节点开启或关闭ssh服务

        Args:
            node_name(str): 节点名
            start(bool): True为开启服务，False为关闭服务
            passwd(str): 为容器的root用户配置密码，初始为123456

        Returns:
            返回True为成功更改服务状态，否则直接报错
        """
        print("Starting SSH service is time-consuming, please be patient...")
        resp = self._post(f"/master/ssh_service/", \
            json={"user": self.user, "topo": self.project, "ne": node_name, "ssh": start, "passwd": passwd})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)
        return True

    def get_port_mapping(self, node_name):
        """获得该项目下一个节点的端口映射，另外包括了节点所在worker的ip地址

        Args:
            node_name(str): 节点名

        Returns:
            节点的端口映射和worker的ip地址
        """
        resp = self._get(f"/master/ssh_service/", \
            json={"user": self.user, "topo": self.project, "ne": node_name})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)
        return {"worker_ip": resp_json["worker_ip"], "ne_port": resp_json["ne_port"]}

    def modify_port_mapping(self, node_name, port_mapping):
        """修改该项目下一个节点（容器）到宿主机的端口映射

        Args:
            node_name(str): 节点名
            port_mapping(list): 端口映射，为一个包含偶数个元素的列表，
                index为偶数为容器端口，奇数为宿主机端口

        Returns:
            返回True为成功更改端口映射状态，否则直接报错
        """
        resp = self._put(f"/master/modify_port_mapping/", \
            json={"user": self.user, "topo": self.project, "ne": node_name, "port_mapping": port_mapping})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)
        return True

    def get_nic_nickname2realname(self, node_name):
        """获取某一节点上所有端口的昵称到真实名字的对应关系

        Args:
            node_name(str): 节点名

        Returns:
            包含节点上所有端口的昵称到真实名字的对应关系字典，否则直接报错
            字典格式为：{'s1h1': 'ea857d34e8', 's1h2': 'b0a7d30a46'} 
        """
        resp = self._get(f"/my/edit/", params={"username": self.user, "toponame": self.project})
        resp_json = self._parse_resp(resp)
        nic_data = resp_json["static"]  # 从相应里提取对应关系信息
        return nic_data[node_name]

    def get_nic_realname2nickname(self, node_name):
        """获取某一节点上所有端口的真实名字到昵称的对应关系

        Args:
            node_name(str): 节点名

        Returns:
            包含节点上所有端口到网卡的对应关系字典，否则直接报错
            字典格式为：{'ea857d34e8': 's1h1', 'b0a7d30a46': 's1h2'}
        """
        nickname2realname_dict = self.get_nic_nickname2realname(node_name)
        return dict([val, key] for key, val in nickname2realname_dict.items())

    def get_node_worker_ip(self, node_name=None):
        """获取某一节点所在的worker的IP地址
        
        若node_name = None，则返回所有节点所在的worker的ip地址

        Args:
            node_name(str): 节点名

        Returns:
            节点所在的worker的ip地址，如：
            {"h1":"192.168.1.1", "h2": "192.168.1.2"}    
        """
        
        resp = self._get(f"/re/project/{self.project}/worker_ip/",
            params={"user": self.user})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)
        all_node_worker_ip_info = resp_json["worker_ip"]
        if node_name == None:
            return all_node_worker_ip_info
        else:
            try:
                one_node_worker_ip_info={}
                one_node_worker_ip_info.setdefault(node_name,
                                                all_node_worker_ip_info[node_name])
                return one_node_worker_ip_info
            except KeyError:
                raise NodeNotExistsError(f"Node [{node_name}] does not exist, "
                    f"avaliable nodes are {list(all_node_worker_ip_info.keys())}")
klonet_api/node.py
Lizonghang/KlonetAI
import time
from .common import Manager, Topo

class ProjectManager(Manager):
    '''项目管理类

    用于项目创建、查询与删除。

    Attributes:
        user(str): 用户名
        project(str): 项目名

    '''
    def __init__(self, user_name, backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name

    def deploy(self, project_name, topo, quiet=False, timeout_min=30,
        pool_interval_s=1):
        '''创建项目，即向后台创建拓扑。

        拓扑的创建意味着一个项目的建立。创建项目的过程为：向后台发送异步拓扑创建请求，
        此时后端会立即返回请求结果；之后，本方法会持续请求进度条API并默认打印进度，
        直至进度为100%，则认为拓扑创建成功。本拓扑创建方法为推荐方法。

        Args:
            project_name(str): 项目名
            topo(Topo): Topo对象。请注意在传入Topo对象前使用Topo对象的add_node和add_link
                方法设计拓扑
            quiet(bool): 默认为False。若为False，则将打印进度；否则将关闭打印
            timeout_min(int): 超时时间（分钟）。默认为30分钟
            pool_interval_s(int): 轮询进度条API的间隔（秒）

        
        '''
        start_time = time.time()
        duration_s = 0

        self.async_deploy(project_name, topo)
        
        while duration_s <= timeout_min * 60:
            progress_value = self._get_progress(project_name, usage="deploy")
            if not quiet:
                print(f"Deployment progress: {progress_value} %")
            if progress_value == 100:
                break
            duration_s = time.time() - start_time
            time.sleep(pool_interval_s)

    def destroy(self, project_name, quiet=False, timeout_min=30,
            pool_interval_s=1):
        '''删除项目

        包含拓扑、网络实验监控服务、流量服务等该项目相关的所有内容。
        删除项目的过程为：向后台发送异步拓扑删除请求，此时后端会立即返回请求结果；之后，
        本方法会持续请求进度条API并默认打印进度，直至进度为100%，则认为项目删除成功。
        本拓扑删除方法为推荐方法。

        Args:
            project_name(str): 项目名
            quiet(bool): 默认为False。若为False，则将打印进度；否则将关闭打印。
            timeout_min(int): 超时时间（分钟）。默认为30分钟。
            pool_interval_s(int): 轮询进度条API的间隔（秒）

        Returns:
            None
        '''
        start_time = time.time()
        duration_s = 0

        self.async_destroy(project_name)
        
        while duration_s <= timeout_min * 60:
            progress_value = self._get_progress(project_name, usage="delete")
            if not quiet:
                print(f"Destruction progress: {progress_value} %")
            if progress_value == 100:
                break
            duration_s = time.time() - start_time
            time.sleep(pool_interval_s)

    def async_deploy(self, project_name, topo):
        '''向后台发送异步拓扑创建请求，令后台开始创建拓扑。

        Args:
            project_name(str): 项目名
            topo(Topo): Topo对象。请注意在传入Topo对象前使用Topo对象的add_node和add_link
                方法设计拓扑

        Returns:
            None
        '''
        payload = {"user": self.user, "topo": project_name,
            "networks": topo.dictform()}
        resp = self._post("/master/topo/", json=payload)
        self._check_resp_code(self._parse_resp(resp))

    def async_destroy(self, project_name):
        '''向后台发送异步拓扑删除请求，令后台开始删除拓扑。

        Args:
            project_name(str): 项目名

        Returns:
            None
        '''
        payload = {"user": self.user, "topo": project_name}
        resp = self._delete("/master/topo/", json=payload)
        self._check_resp_code(self._parse_resp(resp))

    def _get_progress(self, project_name, usage="deploy"):
        '''请求后端进度条API。

        Args:
            project_name(str): 项目名
            usage(str): 进度条类型（如deploy表示“拓扑创建”的进度条）

        Returns:
            一个float类型的变量，其值代表了进度值。100代表进度为100%
        '''
        payload = {"user": self.user, "topo": project_name, "usage": usage}
        resp = self._post("/master/process_bar/", json=payload)
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        return resp_json["process_value"]

    def get_topo(self, project_name):
        '''获取目标已创建项目的Topo对象。

        Args:
            project_name(str): 项目名

        Returns:
            Topo对象
        '''
        resp = self._get(f"/re/project/{project_name}/",
            params={"user": self.user})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        return Topo(**resp_json["project"]["topo"])

    def deploy_with_topo_description_dict(self, project_name,
                                        topo_description_dict):
        '''使用拓扑描述字典来创建项目（不推荐）。

        此API使得用户可通过其它方式生成拓扑描述，
        再调用此API来创建项目。（拓扑描述字典实质上与Topo对象的__dict__相同。）

        Args:
            project_name(str)：项目名
            topo_description_dict(dict)：拓扑描述字典

        '''
        topo = Topo(**topo_description_dict)
        self.deploy(project_name, topo)
    
    def get_projects(self):
        '''获取用户已创建的项目列表

        Returns:
            一个列表，包含了所有已创建的项目名。如：

            ["project1", "project2", "project3", ...]
        '''
        resp = self._get(f"/re/project/",
            params={"user": self.user})
        resp_json = self._parse_resp(resp)
        self._check_resp_code(resp_json)

        return resp_json["topo_list"]
klonet_api/project.py
Lizonghang/KlonetAI
import json
from .common import Traffic, Manager, FlowGeneratorNotExitsError, FlowAPISchemaError, PktLengthToBigPktGen1Error
from ...webserver.schema.schema import parameter_check
from ...webserver.schema.flow_scheme import scheme_pkt_gen1, scheme_pkt_gen2, scheme_traffic_gens
from ...webserver.schema.schema import parameter_check
# from .flow_scheme import scheme_pkt_gen1, scheme_pkt_gen2, scheme_traffic_gen


# from jsonschema import validate, draft7_format_checker
# from jsonschema.exceptions import SchemaError, ValidationError

# def parameter_check(data,schema_data):
#     try:
#         validate(instance=data, schema=schema_data, format_checker=draft7_format_checker)
#     except SchemaError as e:
#         #print("验证模式schema出错:\n出错位置：{}\n提示信息：{}".format(e.path, e.message))
#         raise ValueError("验证模式schema出错:\n出错位置：{}\n提示信息：{}".format(e.path, e.message))
#     except ValidationError as e:
#         print("json数据不符合schema规定:\n出错字段：{}\n提示信息：{}".format(e.path, e.message))
#         return  {'code': 0, 'msg': "您填写的数据不符合规则，请修改! 错误信息：{}".format(e.message)}
#     else:
#         return  {'code': 1, 'msg': "schema参数检查通过！"}


class TrafficEvent():
    '''流量事件类

    用于给管理流量事件，一个流量事件中可能包含多种流

    Attributes:
        traffic_event_name(str): 要添加的流事件的名称，如f1

    '''
    def __init__(self, traffic_event_name) -> None:
        self.traffic_event_name = traffic_event_name
        self.traffic_event = {traffic_event_name:{'traffic_gen':[],
                            'pkt_gen1':[], 'pkt_gen2':[], 'trace':[]}}
        self.traffic_gen_cli_param_dict = {"EARB":"b", "RT":"t", "RN":"n", "SEED":"s"}
    
    def add_flow(self, traffic_generator, **kwargs):
        '''外部调用TrafficEvent添加不同流模式的流

        注意：目前只有三种generator可以指定

        Args:
            traffic_generator(str): 添加的流的generator
            kwargs(dict): 相关参数
        
        Returns:
            None

        Raises:
            FlowGeneratorNotExitsError: 当指定了不存在的generator时触发
        '''
        if traffic_generator == "traffic_gen":
            self.flow_add_traffic_gen(**kwargs)
        elif traffic_generator == "pkt_gen1":
            self.flow_add_pkt_gen1(**kwargs)
        elif traffic_generator == "pkt_gen2":
            self.flow_add_pkt_gen2(**kwargs)
        else:
            raise FlowGeneratorNotExitsError("请指定存在的流量generator(traffic_gen\pkt_gen1\pkt_gen2)")

    def flow_add_traffic_gen(self, **kwargs):
        '''在此流量事件中添加traffic_gen模式的流

        注意：

        Args:
            kwargs(dict): 相关参数
        
        Returns:
            None
        
        Raises:
            FlowAPISchemaError: 参数有误触发此Error
        '''
        para_check_result = parameter_check(kwargs, scheme_traffic_gen)
        if para_check_result['code'] == 1:
            flow_conf = self.get_default_para_traffic_gen(kwargs['server_list'], kwargs['client_name'])
            for key in ["req_size_dist", "dscp", "rate"]:
                if key in kwargs:
                    flow_conf["client"]["client_config"][key] = kwargs[key]
            for key in ["EARB", "RT", "RN", "SEED"]: #期望平均接收带宽、请求时间、请求个数、生成随机数的种子
                if key in kwargs:
                    flow_conf["client"]["cli_param"][key] = kwargs[key]
            self.traffic_event[self.traffic_event_name]['traffic_gen'].append(flow_conf)
        else:
            raise FlowAPISchemaError("[traffic_gen]参数有误，请参照API编程指南中流相关API确保参数正确")


    def flow_add_pkt_gen1(self, **kwargs):
        '''在此流量事件中添加pkt_gen1模式的流

        注意：

        Args:
            kwargs(dict): 相关参数
        
        Returns:
            None
        
        Raises:
            FlowAPISchemaError: 参数有误触发此Error
        '''
        para_check_result = parameter_check(kwargs, scheme_pkt_gen1)
        if para_check_result['code'] == 1:
            if "pkt_length" in kwargs:
                if kwargs["pkt_length"] > "1500":
                    raise PktLengthToBigPktGen1Error(f'''指定pkt_length为{kwargs["pkt_length"]},\
                    不满足小于1500要求''')
            if "ip_id" in kwargs:
                if kwargs["ip_id"] > "65536":
                    raise IOError("ip_id超出范围")
            flow_conf = self.get_default_para_pkt_gen12(if_pkt_gen1=True)
            for key in kwargs:
                flow_conf[key] = kwargs[key]
            self.traffic_event[self.traffic_event_name]['pkt_gen1'].append(flow_conf)
        else:
            raise FlowAPISchemaError("[pkt_gen1]参数有误，请参照API编程指南中流相关API确保参数正确")



    def flow_add_pkt_gen2(self, **kwargs):
        '''在此流量事件中添加pkt_gen2模式的流

        注意：

        Args:
            kwargs(dict): 相关参数
        
        Returns:
            None
        
        Raises:
            FlowAPISchemaError: 参数有误触发此Error
        '''
        para_check_result = parameter_check(kwargs, scheme_pkt_gen2)
        if para_check_result['code'] == 1:
            flow_conf = self.get_default_para_pkt_gen12(if_pkt_gen1=False)
            for key in kwargs:
                flow_conf[key] = kwargs[key]
            self.traffic_event[self.traffic_event_name]['pkt_gen2'].append(flow_conf)
        else:
            raise FlowAPISchemaError("[pkt_gen2]参数有误，请参照API编程指南中流相关API确保参数正确")

    def get_default_para_pkt_gen12(self, if_pkt_gen1):
        '''定义pkt_gen1或pkt_gen2的默认参数

        Args:
            if_pkt_gen1(Boolean): 是否为pkt_gen1
            src(str): 源节点
            dst(str): 目的节点
            src_ip(str): 源IP
            dst_ip(str): 目的IP
        
        Returns:
            pkt_gen1或者pkt_gen2的默认参数
        
        Raises:
            None
        '''
        flow = {
            "rate": "1",
            "duration": "10",
        }
        if if_pkt_gen1:
            pkt_gen1_para = {
                "pkt_length": "1000",
                "dist": "normal",
                "normal_scale": "0.1",
                "ip_tos": "0",
                "ip_ttl": "64",
                "ip_id": "1",
                "proto": "tcp",
                "tcp_header": {
                    "tcp_window": "1000",
                    "sport": "10000",
                    "dport": "10000"
                },
                "udp_header": {}
            }
            flow.update(pkt_gen1_para)
        else:
            pkt_gen2_para = {
                "pkt_length": {
					"1000": "1"
				},
				"on_k": "2",
				"on_min": "1",
				"off_k": "2",
				"off_min": "1"
            }
            flow.update(pkt_gen2_para)
        return flow

    def get_default_para_traffic_gen(self, server_list, client_name):
        '''定义traffic_gen的默认参数

        Args:
            server_list(list): [目的网元的主机名＋端口号]构成的数组
            client_name(str): 客户端网元名，发送网元的主机名
        
        Returns:
            traffic_gen的默认参数
        
        Raises:
            None
        '''
        flow = {
            "mode":"0", #默认为client模式非incast模式
            "server_list": server_list,
            "client": {
                "client_name": client_name,
                "client_config": {
                    "server_list": server_list,
                    "req_size_dist": {
                        "100": "0.1",
                        "200": "0.4",
                        "300": "1"
                    },
                    "dscp": {
                        "0": "50",
                        "1": "50"
                    },
                    "rate": {
                        "1Mbps": "50",
                        "2Mbps": "50"
                    }
				},
                "cli_param": {
                    "b": "1.5",
                    "t": "10",
                    "n": "",
                    "s": "1"
                }
            }
        }
        return flow
    

class TrafficManager(Manager):
    '''流量管理类
    
    用于已创建项目中的流量相关管理
    
    Attributes:
        user_name(str): 用户名
        project_name(str): 实验名
        backend_ip(str): 后端master IP
        backend_port(str): 后端master Port
    
    '''

    def __init__(self, user_name, project_name,
                backend_ip=None, backend_port=None):
        super().__init__(backend_ip, backend_port)
        self.user = user_name
        self.project = project_name
        self.traffics = {
            "user": self.user,
            "traffics":{}
        }

    def add_event(self, trafficEvent:TrafficEvent):
        '''将流数据存储到数据库中

        注意：该API仅对已创建的项目生效！

        Args:
            flow_generator(str): 流量发生器类型(pkt_gen1,pkt_gen2,traffic_gen)
            flow_name(str): 流的名称
            **kwargs(dict): 参数字典，不同流量发生器参数字典不同,详情请看操作手册

        Returns:
            None

        Raises:
            pass
        
        '''
        self.traffics['traffics'].update(trafficEvent.traffic_event)

        # print(payload)
    def save_event(self):
        """
        The save_event function saves the traffic events to the database.
        It takes no arguments and returns nothing.
        
        Args:
            self: Refer to the object of the class
        
        Returns:
            A dictionary with the key &quot;code&quot; and value 0
        
        Doc Author:
            Trelent
        """
        if self.traffics['traffics'] == {}:
            raise RuntimeError("请先使用add_event函数，添加流量事件后再存储")
        resp = self._post(f"/re/project/{self.project}/traffic_app/", self.traffics)
        self._check_resp_code(self._parse_resp(resp))

    def deploy_event(self, traffic_event_name):
        '''将已保存的流事件部署下去

        注意：需要先存储流事件，再执行部署操作

        Args:
            traffic_event_name(str): 流量事件名称

        Returns:
            None

        Raises:
            None
        '''
        payload = {"user":self.user, "topo":self.project, "app_name":traffic_event_name}
        resp = self._post("/master/traffic/", payload)
        self._check_resp_code(self._parse_resp(resp))

    def stop_event(self, traffic_event_name):
        """
        The stop_event function is used to stop a traffic event.
        
        Args:
            traffic_event_name(str): Specify the name of the traffic event that is to be stopped
        
        Returns:
            The response code
        
        Doc Author:
            Trelent
        """
        payload = {"user":self.user, "topo":self.project, "app_name":traffic_event_name}
        resp = self._delete("/master/traffic/", payload)
        self._check_resp_code(self._parse_resp(resp))

    
    def rm_event(self, traffic_event_name, force=False):
        payload = {"user":self.user}
        if force:
            self.stop_event(traffic_event_name)
        resp = self._delete(f"/re/project/{self.project}/traffic_app/{traffic_event_name}/", payload)
        self._check_resp_code(self._parse_resp(resp))

    # def get_flow_monitor_status(self):

        


klonet_api/traffic.py
Lizonghang/KlonetAI
import requests
import klonet_api
from klonet_api import *


def error_handler(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as msg:
            print(f"{msg}")

    return wrapper


def http_response_handler(response, func=lambda r: r["msg"]):
    if response.status_code != 200:
        err_msg = f"Request failed with status code {response.status_code}"
        return err_msg
    else:
        data = response.json()
        code = data.get("code", 1)
        msg = data.get("msg", "Unknown")
        if code == 0:
            return f"Request failed. Error message: {msg}"
        else:
            return func(data)


class KlonetAI:

    def __init__(self):
        self._agent = None
        self._backend_host = ""
        self._port = 0
        self._project = ""
        self._user = ""
        self._image_manager = None
        self._project_manager = None
        self._node_manager = None
        self._link_manager = None
        self._cmd_manager = None
        self._logged_in = False
        self._topo = Topo()
        self._link_config = {}
        self.additional_info = {}

    @property
    def project_name(self):
        return self._project

    @property
    def user(self):
        return self._user

    @property
    def backend_host(self):
        return self._backend_host

    @property
    def port(self):
        return self._port

    @property
    def is_logged_in(self):
        return self._logged_in

    @property
    def is_topo_deployed(self):
        response = self.remote_topo
        return False if type(response) is str else True

    @property
    def is_agent_initialized(self):
        return self._agent is not None

    @property
    def images(self):
        return self._image_manager.get_images()

    @property
    def topo(self):
        return self._topo

    @property
    def remote_topo(self):
        ip = self._backend_host
        port = self._port
        project = self._project
        user = self._user
        url = f"http://{ip}:{port}/re/project/{project}/?user={user}"
        response = requests.get(url)

        def get_topo(data_json):
            project_info = data_json.get("project", {})
            topo = project_info.get("topo", "No topo data")
            return topo
        return http_response_handler(response, get_topo)

    @property
    def nodes(self):
        return self._topo.get_nodes()

    @property
    def remote_nodes(self):
        ip = self._backend_host
        port = self._port
        project = self._project
        user = self._user
        url = f"http://{ip}:{port}/re/project/{project}/node/?user={user}"
        response = requests.get(url)

        def get_node_info(data_json):
            node_info = data_json.get("node_info", {})
            return node_info
        return http_response_handler(response, get_node_info)

    @property
    def links(self):
        return self._topo.get_links()

    @property
    def remote_links(self):
        ip = self._backend_host
        port = self._port
        project = self._project
        user = self._user
        url = f"http://{ip}:{port}/re/project/{project}/link/?user={user}"
        response = requests.get(url)

        def get_link_info(data_json):
            link_info = data_json.get("link_info", {})
            return link_info
        return http_response_handler(response, get_link_info)

    def klonet_login(self, project_name, user_name, host_ip, port):
        self._project = project_name
        self._user = user_name
        self._backend_host = host_ip
        self._port = port
        self._image_manager = ImageManager(self._user, self._backend_host, self._port)
        self._project_manager = ProjectManager(self._user, self._backend_host, self._port)
        self._node_manager = NodeManager(self._user, self._project, self._backend_host, self._port)
        self._link_manager = LinkManager(self._user, self._project, self._backend_host, self._port)
        self._cmd_manager = CmdManager(self._user, self._project, self._backend_host, self._port)
        self._logged_in = True

    def test_klonet_connection(self):
        try:
            _ = self.images
            return True
        except (klonet_api.common.errors.HttpStatusError,
                requests.exceptions.ConnectionError, AttributeError):
            return False

    def create_agent(self, agent=None, agent_name="", key="", tools=[],
                     openai_model="gpt-3.5-turbo-16k", keep_tools=False,
                     chat_prompt_template="LIKirin/klonetai-prompts",
                     run_prompt_template="LIKirin/klonetai-prompts"):
        if agent:
            self._agent = agent
            print(f"Connected to Agent {agent.model_name}.")
        elif agent_name == "openai":
            from agent.openai import KAIOpenAIAgent
            from key import OpenAI_API_Key
            self._agent = KAIOpenAIAgent(
                model=openai_model,
                api_key=key or OpenAI_API_Key,
                chat_prompt_template=chat_prompt_template,
                run_prompt_template=run_prompt_template,
                additional_tools=tools
            )
            print("Connected to OpenAI Agent.")
        else:
            from agent.huggingface import KAIHfAgent
            from key import Huggingface_API_Key
            pinned_model = {
                "starcoder": "bigcode/starcoder",
                "codellama-13b-hf": "codellama/CodeLlama-13b-hf",
                "codeLlama-34b-instruct-hf": "codellama/CodeLlama-34b-Instruct-hf",
                "phind-codellama-34b-v2": "Phind/Phind-CodeLlama-34B-v2",
            }
            model_suffix = pinned_model.get(agent_name, "bigcode/starcoder")
            url = f"https://api-inference.huggingface.co/models/{model_suffix}"
            self._agent = KAIHfAgent(
                url,
                token=key or Huggingface_API_Key,
                chat_prompt_template=chat_prompt_template,
                run_prompt_template=run_prompt_template,
                additional_tools=tools
            )
            print(f"Connected to 🤗 Agent: {model_suffix}.")

        # Remove unnecessary built-in tools.
        if not keep_tools:
            self._agent.toolbox.pop("document_qa")
            self._agent.toolbox.pop("image_captioner")
            self._agent.toolbox.pop("image_qa")
            self._agent.toolbox.pop("image_segmenter")
            self._agent.toolbox.pop("transcriber")
            self._agent.toolbox.pop("text_qa")
            self._agent.toolbox.pop("text_classifier")
            self._agent.toolbox.pop("text_reader")
            self._agent.toolbox.pop("translator")
            self._agent.toolbox.pop("image_transformer")
            self._agent.toolbox.pop("image_generator")
            self._agent.toolbox.pop("video_generator")

        # Replace the built-in summarizer with ours.
        import tool
        self._agent.toolbox["summarizer"] = tool.SummarizeTool()
        print(f"All available tools:", list(self._agent.toolbox.keys()))

    @property
    def toolbox(self):
        return self._agent.toolbox

    @property
    def chat_history(self):
        return self._agent.chat_history

    @property
    def chat_state(self):
        return self._agent.chat_state

    @error_handler
    def new_chat(self):
        self._agent.prepare_for_new_chat()

    @error_handler
    def chat(self, *args, **kwargs):
        return self._agent.chat(*args, **kwargs)

    @error_handler
    def run(self, *args, **kwargs):
        return self._agent.run(*args, **kwargs)

    @property
    def mode(self):
        return self._agent.mode

    def set_mode(self, mode):
        self._agent.set_mode(mode)

    def reset_project(self):
        self._topo = Topo()
        self._link_config.clear()
        self._project_manager.destroy(self._project)

    def add_node(self, name, image, cpu_limit=None, mem_limit=None, x=0, y=0):
        node = self._topo.add_node(
            image, name,
            resource_limit={"cpu": cpu_limit, "mem": mem_limit},
            location={"x": x, "y": y})
        return node

    def add_node_runtime(self, name, image, cpu_limit=None, mem_limit=None, x=0, y=0):
        node = self._node_manager.dynamic_add_node(
            name, image,
            resource_limit={"cpu": cpu_limit, "mem": mem_limit},
            location={"x": x, "y": y}
        )
        return node

    def delete_node_runtime(self, name):
        self._node_manager.dynamic_delete_node(name)

    def add_link(self, src_node, dst_node, link_name=None, src_ip="", dst_ip=""):
        link = self._topo.add_link(
            src_node, dst_node, link_name, src_ip, dst_ip)
        return link

    def add_link_runtime(self, src_node, dst_node, link_name=None, src_ip="", dst_ip=""):
        self._link_manager.dynamic_add_link(
            link_name, src_node, dst_node, src_ip, dst_ip)

    def delete_link_runtime(self, link_name):
        self._link_manager.dynamic_delete_link(link_name)

    def configure_link(self, config):
        link_name = config["link"]
        _ = self._link_config.setdefault(link_name, {})

        node_name = config["ne"]
        src_node = self.links[link_name].source
        dst_node = self.links[link_name].target
        src_link_config = _.setdefault(src_node, {})
        dst_link_config = _.setdefault(dst_node, {})

        if node_name == src_node:
            src_link_config.update(**config)
            dst_link_config.update({"link": link_name, "ne": dst_node})
        elif node_name == dst_node:
            src_link_config.update({"link": link_name, "ne": src_node})
            dst_link_config.update(**config)
        else:
            raise LinkInconsistentError(f"Node {node_name} is not on link {link_name}.")

        src_config_obj = LinkConfiguration(**src_link_config)
        dst_config_obj = LinkConfiguration(**dst_link_config)
        self._link_manager.config_link(src_config_obj, dst_config_obj)
        return src_link_config if node_name == src_node else dst_link_config

    def reset_link(self, link_name, clean_cache=False):
        self._link_manager.clear_link_configuration(link_name)
        if clean_cache: self._link_config.clear()

    def query_link(self, link_name, node_name):
        ip = self._backend_host
        port = self._port
        url = f"http://{ip}:{port}/master/linkquery/"
        data = {
            "user": self._user,
            "topo": self._project,
            "links": [{
                "link": link_name,
                "ne": node_name
            }]
        }
        response = requests.post(url, json=data)

        def get_link_info(data_json):
            return data_json["static"]
        return http_response_handler(response, get_link_info)

    def deploy(self):
        self._project_manager.deploy(self._project, self._topo)

    def check_deployed(self):
        ip = self._backend_host
        port = self._port
        user = self._user
        project = self._project
        url = f"http://{ip}:{port}/master/topo/?user={user}&topo={project}"
        response = requests.get(url)

        def get_deploy_status(data_json):
            return data_json["stat"]
        return http_response_handler(response, get_deploy_status)

    def execute(self, node_name, command, block="false", timeout=60):
        response = self._cmd_manager.exec_cmds_in_nodes({
            node_name: [command]
        }, block, timeout)
        return response

    def batch_exec(self, ctns, command, block="false", timeout=60):
        url = f"http://{self._backend_host}:{self._port}/master/batch_exec_cmd/"
        data = {
            "user": self._user,
            "topo": self._project,
            "ctns": ctns,
            "cmd": command,
            "block": block,
            "cmd_timeout_s": timeout
        }
        response = requests.post(url, json=data)

        def get_exec_result(data_json):
            return data_json["exec_results"]
        return http_response_handler(response, get_exec_result)

    def enable_ssh_service(self, node_name):
        return self._node_manager.ssh_service(node_name, True)

    def port_mapping(self, node_name, container_port, host_port):
        return self._node_manager.modify_port_mapping(
            node_name, [container_port, host_port])

    def get_port_mapping(self, node_name):
        return self._node_manager.get_port_mapping(node_name)

    def get_worker_id(self, node_name=None):
        return self._node_manager.get_node_worker_ip(node_name)

    def deploy_from_config(self, config):
        self._topo = Topo(**config)
        url = f"http://{self._backend_host}:{self._port}/master/topo/"
        data = {
            "user": self._user,
            "topo": self._project,
            "networks": config
        }
        response = requests.post(url, json=data)
        return http_response_handler(response)

    def create_template_topo(self, config):
        url = f"http://{self._backend_host}:{self._port}/generate"
        response = requests.post(url, json=config)

        def get_topo_config(data_json):
            return data_json["net"]
        return http_response_handler(response, get_topo_config)

    def config_public_network(self, node_name, turn_on=True):
        url = f"http://{self._backend_host}:{self._port}/master/node/network/"
        data = {
            "user": self._user,
            "topo": self._project,
            "ne": node_name
        }
        req_func = requests.post if turn_on else requests.delete
        response = req_func(url, json=data)
        return http_response_handler(response)

    def check_public_network(self, node_name):
        ip = self._backend_host
        port = self._port
        user = self._user
        project = self._project
        url = f"http://{ip}:{port}/master/node/network/?user={user}&topo={project}&ne={node_name}"
        response = requests.get(url)

        def get_status(data_json):
            return data_json["status"]
        return http_response_handler(response, get_status)

    def upload_file(self, node_name, src_file, tgt_filepath="/home"):
        url = f"http://{self._backend_host}:{self._port}/file/uload/"
        data = {
            "user": self._user,
            "topo": self._project,
            "ne_name": node_name,
            "file_path": tgt_filepath,
        }
        files = {"file": open(src_file, "rb") if type(src_file) is str else src_file}
        response = requests.post(url, data=data, files=files)
        return http_response_handler(response)

    def manage_worker(self, worker_ip, delete_worker=False):
        url = f"http://{self._backend_host}:{self._port}/master/worker/{worker_ip}/"
        data = {"worker_ip": worker_ip}
        req_func = requests.delete if delete_worker else requests.post
        response = req_func(url, json=data)
        return http_response_handler(response)

    def check_health(self):
        ip = self._backend_host
        port = self._port
        user = self._user
        project = self._project
        url = f"http://{ip}:{port}/master/heartbeat_health/?user={user}&project={project}"
        response = requests.get(url)

        def get_broken_nodes(data_json):
            is_broken = data_json["is_broken"]
            broken_nodes = data_json["broken_nes"]
            return is_broken, broken_nodes
        return http_response_handler(response, get_broken_nodes)

klonetai.py
Lizonghang/KlonetAI
from .qa import template as qa_template

template/__init__.py
Lizonghang/KlonetAI
from string import Template

template = Template('''

**Instruction Template for Question-Answer ChatGPT Application**

This template is designed to provide answers to questions based on 
a given context text and a query. Please follow the guidelines below 
for effective interactions:

===

1. **Context Text:**
    - Provide a clear and concise context text that contains information relevant to the query.
    - The context text is wrapped by <begin> <end>, e.g., <begin> I am Lucy and live in London. <end>

Context: <begin> $context <end>

===

2. **Query:**
    - Be as concise as possible while providing enough information for a meaningful response.
    - The query text is wrapped by <begin> <end>, e.g., <begin> What is my name? <end>

Query: <begin> $query <end>

===

3. **Response:**
    - The response will be provided in code or text form, depending on the goal of the query.
    - If requesting code, specify the programming language or platform you want the code in.
        default to Python.
''')

template/qa.py
Lizonghang/KlonetAI
from .common import (
    GetCurrentDateTimeTool,
    GetCurrentDayTool
)
from .klonet import (
    KlonetAddLinkTool,
    KlonetAddNodeTool,
    KlonetCommandExecTool,
    KlonetRuntimeDeleteNodeTool,
    KlonetDeployTool,
    KlonetCheckDeployedTool,
    KlonetGetAllImagesTool,
    KlonetViewTopoTool,
    KlonetDestroyProjectTool,
    KlonetRuntimeAddNodeTool,
    KlonetRuntimeAddLinkTool,
    KlonetRuntimeDeleteLinkTool,
    KlonetSSHServiceTool,
    KlonetPortMappingTool,
    KlonetGetPortMappingTool,
    KlonetGetIPTool,
    KlonetLinkConfigurationTool,
    KlonetResetLinkConfigurationTool,
    KlonetLinkQueryTool,
    KlonetGetWorkerIPTool,
    KlonetTreeTopoTemplate,
    KlonetStarTopoTemplate,
    KlonetFatTreeTopoTemplate,
    KlonetLinearTopoTemplate,
    KlonetConfigurePublicNetworkTool,
    KlonetCheckPublicNetworkTool,
    KlonetFileDownloadTool,
    KlonetFileUploadTool,
    KlonetManageWorkerTool,
    KlonetCheckHealthTool,
)
from .gpt import SummarizeTool

base = (
    GetCurrentDateTimeTool,
    GetCurrentDayTool,
)

klonet = (
    KlonetAddLinkTool,
    KlonetAddNodeTool,
    KlonetCommandExecTool,
    KlonetRuntimeDeleteNodeTool,
    KlonetDeployTool,
    KlonetCheckDeployedTool,
    KlonetGetAllImagesTool,
    KlonetViewTopoTool,
    KlonetDestroyProjectTool,
    KlonetRuntimeAddNodeTool,
    KlonetRuntimeAddLinkTool,
    KlonetRuntimeDeleteLinkTool,
    KlonetSSHServiceTool,
    KlonetPortMappingTool,
    KlonetGetPortMappingTool,
    KlonetGetIPTool,
    KlonetLinkConfigurationTool,
    KlonetResetLinkConfigurationTool,
    KlonetLinkQueryTool,
    KlonetGetWorkerIPTool,
    KlonetConfigurePublicNetworkTool,
    KlonetCheckPublicNetworkTool,
    KlonetFileDownloadTool,
    KlonetFileUploadTool,
    KlonetManageWorkerTool,  # TODO: To be test.
    KlonetCheckHealthTool,
)

topo = (
    KlonetTreeTopoTemplate,
    KlonetStarTopoTemplate,
    KlonetFatTreeTopoTemplate,
    KlonetLinearTopoTemplate,
)

gpt = (
    SummarizeTool,
)

tool/__init__.py
Lizonghang/KlonetAI
import datetime
from transformers import Tool


class GetCurrentDateTimeTool(Tool):
    name = "get_current_datetime"
    description = ('''
    Get the current date and time.

    Args:
        None

    Returns:
        str: A string representing the current date and time in the 
            "YYYY-MM-DD HH:MM:SS" format.

    Example:
        >>> time = get_current_datetime()
        >>> print(time)
    ''')

    outputs = ["str"]

    def __call__(self):
        date_obj = datetime.datetime.now()
        date_str = date_obj.strftime("%Y-%m-%d %H:%M:%S")
        return date_str


class GetCurrentDayTool(Tool):
    name = "get_current_day"
    description = ('''
    Get the current day of week.

    Args:
        None

    Returns:
        str: A string representing the day of the week.

    Example:
        >>> day = get_current_day()
        >>> print(day)
    ''')

    outputs = ["str"]

    def __call__(self):
        date_obj = datetime.datetime.now()
        day_str = date_obj.strftime("%A")
        return day_str

tool/common.py
Lizonghang/KlonetAI
import openai
from transformers import Tool
from key import OpenAI_API_Key
from template import qa_template

openai.api_key = OpenAI_API_Key


class SummarizeTool(Tool):
    name = "summarize_query"
    description = ('''
    Summarize a long tutorial text into a shorter text based on a query. It 
    aims to reduce the subsequent API call costs by providing a concise summary.
    
    Args:
        context (str): The long tutorial text to be summarized.
        query (str): The query that guides the summarization process.
        
    Returns:
        str: A concise summary of the tutorial text based on the query. It is 
            automatically determined by the agent based on the prompt's needs.
            
    
    Example:
        # Replace <Your-URL-Path> to the document you need.
        >>> tutorial = text_downloader(url="<Your-URL-Path>")
        # Replace <Your-Query-Text> to the query you need.
        >>> user_query = "<Your-Query-Text>"
        >>> summary = summarize_query(context=tutorial, query=user_query)
    ''')

    inputs = ["str", "str"]
    outputs = ["str"]

    def __call__(self, context: str, query: str):
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo-16k",
            messages=[{
                "role": "user",
                "content": qa_template.substitute(
                    context=context,
                    query=query
                )
            }]
        )
        content = response['choices'][0]['message']['content']
        context = content.replace('```', '\n')
        return context

tool/gpt.py
Lizonghang/KlonetAI
from transformers import Tool
from klonetai import KlonetAI, error_handler


kai = KlonetAI()


class KlonetGetAllImagesTool(Tool):
    name = "klonet_get_all_images"
    description = ('''
    This method fetches a list of all image names and their associated 
    Docker Image objects.

    Args: 
        None

    Returns:
        None

    Example:
        >>> klonet_get_all_images()
    ''')

    @error_handler
    def __call__(self):
        kai.images.keys()


class KlonetViewTopoTool(Tool):
    name = "klonet_view_topo"
    description = ('''
    Overview the current network topology on Klonet backend.
    
    Example:
        >>> klonet_view_topo()
    ''')

    @error_handler
    def __call__(self):
        result = kai.remote_topo
        print(result)


class KlonetAddNodeTool(Tool):
    name = "klonet_add_node"
    description = ('''
    Add a node to the Klonet network. The node with position (x, y) cannot overlap with
    other nodes.
    
    Args:
        name (str): The name of the node being added. The name of this new node cannot be 
            the same as existing nodes.
        image (str): The name of Docker image used by the node. Use the klonet_get_all_images
            tool to see the available images.
        x (int): The x-coordinate of the node on the canvas. It should be set as an integer 
            value between 0 and 700. The agent should set this for pretty visual layout.
        y (int): The y-coordinate of the node on the canvas. It should be set as an integer 
            value between 0 and 700. The agent should set this for pretty visual layout.
        cpu_limit (int, optional): CPU utilization limit for the node, unit: %, 
            default to None, which will use the default cpu limits from the Docker image.
        mem_limit (int, optional): Memory utilization limit for the node, unit: Mbytes, 
            default to None, which will use the default memory limits from the Docker image.
    
    Returns:
        None

    Raises:
        NodeDuplicatesError: Raised when the provided node name is a duplicate.

    Example:
        # Add a ubuntu host named h1.
        >>> klonet_add_node("h1", "ubuntu", x=500, y=500)
        # Add an OVS switch named s1.
        >>> klonet_add_node("s1", "ovs", x=500, y=500)
    ''')

    inputs = ["str", "str", "int", "int", "int", "int"]

    @error_handler
    def __call__(self, name: str, image: str, x: int, y: int,
                 cpu_limit: int = None, mem_limit: int = None):
        node = kai.add_node(
            name, kai.images[image], cpu_limit, mem_limit, x, y)
        print(f"A new node (name: {node.name}, image: {node.image_name}, "
              f"resource limit: {node.resource_limit}) have been added to the network.")


class KlonetRuntimeAddNodeTool(Tool):
    name = "klonet_runtime_add_node"
    description = ('''
    Add a node to the Klonet network. This tool is designed for post-deployment 
    use, allowing adding nodes to the network during runtime. The node with position 
    (x, y) cannot overlap with other nodes.
    
    Args:
        name (str): The name of the node being added. Be sure not to use the same name as 
            existing nodes.
        image (str): The name of Docker image used by the node. Use the klonet_get_all_images
            tool to see the available images.
        x (int): The x-coordinate of the node on the canvas. It should be set as an integer 
            value between 0 and 700. The agent should set this for pretty visual layout.
        y (int): The y-coordinate of the node on the canvas. It should be set as an integer 
            value between 0 and 700. The agent should set this for pretty visual layout.
        cpu_limit (int, optional): CPU utilization limit for the node, unit: %, 
            default to None, which will use the default cpu limits from the Docker image.
        mem_limit (int, optional): Memory utilization limit for the node, unit: Mbytes, 
            default to None, which will use the default memory limits from the Docker image.

    Returns:
        None

    Raises:
        NodeDuplicatesError: Raised when the provided node name is a duplicate.
    
    Example:
        # Add a ubuntu host named h2.
        >>> klonet_runtime_add_node("h2", "ubuntu", x=500, y=500)
        # Add an OVS switch named s1.
        >>> klonet_runtime_add_node("s2", "ovs", x=500, y=500)
    ''')

    inputs = ["str", "str", "int", "int", "int", "int"]

    @error_handler
    def __call__(self, name: str, image: str, x: int, y: int,
                 cpu_limit: int = None, mem_limit: int = None):
        node = kai.add_node_runtime(
            name, kai.images[image], cpu_limit, mem_limit, x, y)
        print(f"A new node (name: {node.name}, image: {node.image_name}, "
              f"resource limit: {node.resource_limit}) have been added to the network.")


class KlonetRuntimeDeleteNodeTool(Tool):
    name = "klonet_runtime_delete_node"
    description = ('''
    Delete a node from the Klonet network by its name. This tool is designed for 
    post-deployment use, allowing deletion of nodes from the network during runtime.

    Args:
        name (str): The name of the node to delete.

    Returns:
        None

    Example:
        # Replace h1 with the name of the node you want to delete.
        >>> klonet_runtime_delete_node("h1")  
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, name: str):
        kai.delete_node_runtime(name)
        print(f"Node {name} has been removed.")


class KlonetAddLinkTool(Tool):
    name = "klonet_add_link"
    description = ('''
    Add a network link between nodes <src_node> and <dst_node> in the Klonet network.
    
    Args:
        src_node (str): The name of the source node.
        dst_node (str): The name of the destination node.
        link_name (str): The name of the link. The name of this new link cannot be 
            the same as existing links.
        src_ip (str): The source IP address. Avoid using the first two and last IP 
            addresses in the subnet. For example, avoid using 10.0.0.0, 10.0.0.1, 
            and 10.0.0.255 in the subnet 10.0.0.0/24.

    Returns:
        None
    
    Example:
        # Replace "h1", "h2", "s1" with the names of nodes you want to link.
        >>> klonet_add_link("h1", "s1", "l1", "10.0.0.2/24")
        >>> klonet_add_link("h2", "s1", "l2", "10.0.0.3/24")
    ''')

    inputs = ["str", "str", "str", "str"]

    @error_handler
    def __call__(self, src_node: str, dst_node: str, link_name: str, src_ip: str):
        src_node = kai.nodes[src_node]
        dst_node = kai.nodes[dst_node]
        link = kai.add_link(src_node, dst_node, link_name, src_ip)
        print(f"A link with name ({link.name}) was added between nodes "
              f"{link.source} (IP: {link.sourceIP}) and {link.target}")


class KlonetRuntimeAddLinkTool(Tool):
    name = "klonet_runtime_add_link"
    description = ('''
    Add a network link between nodes <src_node> and <dst_node> in the Klonet network.
    This tool is designed for post-deployment use, allowing adding links to the network 
    during runtime. 
    
    Args:
        src_node (str): The name of the source node.
        dst_node (str): The name of the destination node.
        link_name (str): The name of the link. The name of this new link cannot be 
            the same as existing links.
        src_ip (str): The source IP address. Avoid using the first two and last IP 
            addresses in the subnet. For example, avoid using 10.0.0.0, 10.0.0.1, 
            and 10.0.0.255 in the subnet 10.0.0.0/24.
    
    Returns:
        None
    
    Example:
        # Replace "h1", "h2", "s1" with the names of nodes you want to link.
        >>> klonet_add_link("h1", "s1", "l1", "10.0.0.2/24")
        >>> klonet_add_link("h2", "s1", "l2", "10.0.0.3/24")
    ''')

    inputs = ["str", "str", "str", "str"]

    @error_handler
    def __call__(self, src_node: str, dst_node: str, link_name: str, src_ip: str):
        src_node = kai.nodes[src_node]
        dst_node = kai.nodes[dst_node]
        kai.add_link_runtime(src_node, dst_node, link_name, src_ip)
        print(f"A link with name ({link_name}) was added between nodes "
              f"{src_node.name} (IP: {src_ip}) and {dst_node.name}")


class KlonetRuntimeDeleteLinkTool(Tool):
    name = "klonet_runtime_delete_link"
    description = ('''
    Delete a network link by its name. This tool is designed for post-deployment use, 
    allowing deletion of links from the network during runtime. 
    
    Args:
        name (str): The name of the link to delete.

    Returns:
        None

    Example:
        # Replace l1 with the name of the link you want to delete.
        >>> klonet_runtime_delete_link("l1")  
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, name: str):
        kai.delete_link_runtime(name)
        print(f"Link {name} has been removed.")


class KlonetDeployTool(Tool):
    name = "klonet_deploy_network"
    description = ('''
    Deploy the designed network to Klonet.
    
    Inputs:
        None

    Returns:
        None
        
    Example:
        >>> klonet_deploy_network()
    ''')

    @error_handler
    def __call__(self):
        kai.deploy()
        print(f"Deploy project {kai.project_name} success.")


class KlonetCheckDeployedTool(Tool):
    name = "klonet_check_deployed"
    description = ('''
    Check whether the network is deployed on remote Klonet backend. 
    Run this tool to confirm whether some operations can be performed. 
    
    Args:
        None
    
    Returns:
        bool: True if the network has been deployed, False otherwise.
    
    Example:
        >>> is_deployed = klonet_check_deployed()
        >>> print(f"Is this project deployed: {is_deployed}.")
    ''')

    outputs = ["bool"]

    @error_handler
    def __call__(self):
        return kai.check_deployed()


class KlonetDestroyProjectTool(Tool):
    name = "klonet_destroy_project"
    description = ('''
    Destroy the current project in Klonet.
    
    Inputs:
        None
    
    Returns:
        None
    
    Example:
        >>> klonet_destroy_project()
    ''')

    @error_handler
    def __call__(self):
        kai.reset_project()
        print("This project has been deleted.")


class KlonetCommandExecTool(Tool):
    name = "klonet_command_exec"
    description = ('''
    Execute the given command on the specified node.

    Args:
        node_name (str): The name of the node where the command will be executed.
        command (str): The command to execute on the specified node. Multiple 
            commands can be separated by semicolons ';', and they will be executed sequentially.

    Returns: 
        None
    
    Example:
        >>> klonet_command_exec("h1", "ls; pwd")
    ''')

    inputs = ["str", "str"]

    @error_handler
    def __call__(self, node_name: str, command: str):
        full_command = "source ~/.bashrc"
        conda_path = kai.additional_info.get("conda_path", "")
        conda_env = kai.additional_info.get("conda_env", "")
        if conda_path and conda_env:
            full_command = " && ".join([
                full_command,
                f"source {conda_path}",
                f"conda activate {conda_env}"]
            )
        full_command = " && ".join([full_command, command])
        response = kai.execute(node_name, f'bash -c "{full_command}"')
        print(list(response[node_name].values())[0]['output'].strip())


class KlonetBatchCommandExecTool(Tool):
    name = "klonet_batch_command_exec"
    description = ('''
    Execute the given command on batched nodes.
    
    Args:
        node_list (list): A list of target nodes where the command will be executed.
        node_type (str): The type of nodes, could be: "all", "specified_ctn_type", 
            or "specified_ctn_list".
            - When node_type is "all", the command is executed on all nodes, and 
                node_list should be empty.
            - When node_type is "specified_ctn_type", the nodes are selected by their
                type, and the selected nodes should execute the command. The types 
                include: "hosts", "controllers", "routers", and "switches". For example,
                when node_list is ["hosts"], all host nodes should execute the command.
            - When node_type is "specified_ctn_list", the nodes are selected by their
                name, and the selected nodes should execute the command. For example, 
                when node_list is ["h1", "s1"], the nodes h1 and s1 should execute the
                command.
        command (str): The command to be executed on the selected nodes. Multiple 
            commands can be separated by semicolons ';', and they will be executed 
            sequentially.
    
    Returns:
        None
        
    Example:
        # Case 1: Run a command on all nodes.
        >>> klonet_batch_command_exec([], "all", "ifconfig")
        # Case 2: Run a command on all host nodes.
        >>> klonet_batch_command_exec(["hosts"], "specified_ctn_type", "ifconfig")
        # Case 3: Run a command on the selected nodes.
        >>> klonet_batch_command_exec(["h1", "s1"], "specified_ctn_list", "ifconfig")
    ''')

    inputs = ["list", "str", "str"]

    @error_handler
    def __call__(self, node_list: list, node_type: str, command: str):
        ctns = {"list_type": node_type, "list": node_list}

        full_command = "source ~/.bashrc"
        conda_path = kai.additional_info.get("conda_path", "")
        conda_env = kai.additional_info.get("conda_env", "")
        if conda_path and conda_env:
            full_command = " && ".join([
                full_command,
                f"source {conda_path}",
                f"conda activate {conda_env}"]
            )
        full_command = " && ".join([full_command, command])
        result = kai.batch_exec(ctns, f'bash -c "{full_command}"')
        print(result)


class KlonetSSHServiceTool(Tool):
    name = "klonet_enable_ssh_service"
    description = ('''
    Enable SSH service on a specified Klonet node.
    
    Args:
        node_name (str): The name of the node on which SSH service should be enabled.
    
    Returns:
        None
    
    Example:
        >>> klonet_enable_ssh_service('h1')
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, node_name: str):
        success = kai.enable_ssh_service(node_name)
        print(f"SSH service on {node_name} started {'success' if success else 'failed'}.")


class KlonetPortMappingTool(Tool):
    name = "klonet_port_mapping"
    description = ('''
    Perform port mapping for a specified Klonet node.
    
    Args:
        node_name (str): The name of the node on which a port to be mapped.
        container_port (int): The port number in the container to be mapped.
        host_port (int): The port number on the host to map to. The host port
            should be between 9200 and 40000.
    
    Returns:
        None
    
    Example:
        # Map the port 80 of the container to the port 8080 of the host.
        >>> klonet_port_mapping('h1', 80, 8080)
    ''')

    inputs = ["str", "int", "int"]

    @error_handler
    def __call__(self, node_name: str, container_port: int, host_port: int):
        success = kai.port_mapping(node_name, container_port, host_port)
        print(f"Port mapping on {node_name} {'success' if success else 'failed'}.")


class KlonetGetPortMappingTool(Tool):
    name = "klonet_get_port_mapping"
    description = ('''
    Show the port map for a given node. 
    
    Args:
        node_name (str): The name of the node to query.
    
    Returns:
        None
    
    Example:
        >>> klonet_get_port_mapping()
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, node_name: str):
        result = kai.get_port_mapping(node_name)
        print(f"Port map: {result['ne_port']}")


class KlonetGetIPTool(Tool):
    name = "klonet_get_ip"
    description = ('''
    Retrieve the IP address of a Klonet node by its name. When you need 
    to communicate with a target host, use this tool to retrieve the 
    target IP address. For example, if h1 wants to ping h2, it should
    call klonet_get_ip('h2') to retrieve its IP address (e.g., 10.0.0.3),
    and then ping 10.0.0.3.
    
    Args:
        node_name (str): The name of the Klonet node for which you want 
            to retrieve the IP address.
    
    Returns:
        str: The IP address of the specified node.
    
    Example:
        >>> ip_address = klonet_get_ip('h1')
        >>> print(ip_address)
    ''')

    inputs = ["str"]
    outputs = ["str"]

    @error_handler
    def __call__(self, node_name: str):
        return kai.nodes[node_name].interfaces[0]['ip']


class KlonetLinkConfigurationTool(Tool):
    name = "klonet_configure_link"
    description = ('''
    This tool allows you to customize the network link settings for a specific node,
    including bandwidth, delay, jitter, packet loss, and queue size, to simulate 
    various network conditions. 
    
    Args:
        link_name (str): The name of the link to configure.
        node_name (str): The name of the node to which the link is connected.
        bandwidth (int, optional): The desired link bandwidth in kbps (must be a 
            positive value).
        delay (int, optional): The desired link delay in microseconds (must be a 
            non-negative value).
        delay_dist (str, optional): The distribution of delay variations. Options 
            include: uniform, normal, pareto, and paretonomal.
        jitter (int, optional): The desired jitter in microseconds (must be a 
            non-negative value).
        correlation (int, optional): The percentage of jitter correlation (must 
            be a non-negative percentage).
        loss (int, optional): The desired packet loss percentage (must be a 
            non-negative percentage).
        queue_size (int, optional): The desired queue size in bytes (must be a 
            non-negative value).
    
    Returns:
        None
    
    Example:
        >>> klonet_configure_link(
                'l1', 'h1', 
                bandwidth = 1000, 
                delay = 30, 
                delay_dist = 'normal',
                jitter = 100,
                correlation = 1,
                loss = 1,
                queue_size = 200000
            )
    ''')

    inputs = ["str", "str", "int", "int", "str", "int", "int", "int", "int"]

    @error_handler
    def __call__(
        self,
        link_name: str,
        node_name: str,
        bandwidth: int = -1,
        delay: int = -1,
        delay_dist: str = "",
        jitter: int = -1,
        correlation: int = -1,
        loss: int = -1,
        queue_size: int = -1
    ):
        config = {
            "link": link_name,
            "ne": node_name,
            **({"bw_kbps": bandwidth} if bandwidth > 0 else {}),
            **({"delay_us": delay} if delay >= 0 else {}),
            **({"delay_distribution": delay_dist} if delay_dist else {}),
            **({"jitter_us": jitter} if jitter >= 0 else {}),
            **({"correlation": correlation} if correlation >= 0 else {}),
            **({"loss": loss} if loss >= 0 else {}),
            **({"queue_size_bytes": queue_size} if queue_size >= 0 else {}),
        }
        # Reset the link configuration before modifying it.
        kai.reset_link(link_name, clean_cache=False)
        merged_config = kai.configure_link(config)
        print(f"Link {link_name} (on the {node_name} side) is configured "
              f"with: {merged_config}")


class KlonetResetLinkConfigurationTool(Tool):
    name = "klonet_reset_link"
    description = ('''
    Reset the configuration of a given network link.
    
    Args:
        link_name (str): The name of the link to reset.

    Returns:
        None
    
    Example:
        >>> klonet_reset_link('l1')
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, link_name: str):
        kai.reset_link(link_name, clean_cache=True)
        print(f"Link {link_name} has been reset.")


class KlonetLinkQueryTool(Tool):
    name = "klonet_link_query"
    description = ('''
    Query information about a link.
    
    Args:
        link_name (str): The name of the link to query.
        node_name (str): The name of the node associated with the link to query.
    
    Returns:
        dict: A dictionary containing information about the queried link.
    
    Example:
        >>> link_info = klonet_link_query("l1", "h1")
    ''')

    inputs = ["str", "str"]
    outputs = ["dict"]

    def __call__(self, link_name: str, node_name: str):
        return kai.query_link(link_name, node_name)


class KlonetGetWorkerIPTool(Tool):
    name = "klonet_get_worker_ip"
    description = ('''
    Show the worker IP for a given node. The worker IP is the IP address 
    of the host machine where the given node is deployed on.

    Args:
        node_name (str, optional): The name of the node to query. If given,
            query only the given node. If not given, query the worker IP of
            all the nodes.

    Returns:
        None

    Example:
        >>> klonet_get_worker_ip()
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, node_name: str = None):
        result = kai.get_worker_id(node_name)
        print(f"Worker IP: {result}")


class KlonetTreeTopoTemplate(Tool):
    name = "klonet_tree_topo_template"
    description = ('''
    Deploy a tree network topology on Klonet. Do not call the
    klonet_deploy_network tool if you use this template.
    
    Args:
        subnet (str): The subnet to deploy the topology.
        ndepth (int, optional): The depth of the tree (default is 2).
        nbranch (int, optional): The number of branches at each level 
            (default is 2).
        density (int, optional): The number of hosts connected to each 
            leaf switch (default is 2).
    
    Returns:
        None
    
    Example:
        >>> klonet_tree_topo_template("192.168.1.0/24", ndepth=2, nbranch=2, density=1)
    ''')

    inputs = ["str", "int", "int", "int"]

    @error_handler
    def __call__(self, subnet: str, ndepth: int = 2, nbranch: int = 2, density: int = 1):
        print("[Warning] This operation will overwrite the existing topology.")
        kai.reset_project()
        print("[Warning] Creating topology in this way will not layout the view.")
        params_config = {
            "topology_type": "tree",
            "tree_depths": ndepth,
            "tree_branches": nbranch,
            "tree_host_density": density,
            "host_counter": 1,
            "switch_counter": 1,
            "link_counter": 1,
            "ip_prefix": subnet
        }
        net_config = kai.create_template_topo(params_config)
        result = kai.deploy_from_config(net_config)
        print(result)


class KlonetStarTopoTemplate(Tool):
    name = "klonet_star_topo_template"
    description = ('''
    Deploy a star network topology on Klonet. Do not call the
    klonet_deploy_network tool if you use this template.
    
    Args:
        subnet (str): The subnet to deploy the topology.
        nstar (str, optional): The number of host nodes (default is 3).
    
    Returns:
        None
    
    Example:
        >>> klonet_star_topo_template("192.168.1.0/24", nstar=3)
    ''')

    inputs = ["str", "int"]

    @error_handler
    def __call__(self, subnet: str, nstar: int = 3):
        print("[Warning] This operation will overwrite the existing topology.")
        kai.reset_project()
        print("[Warning] Creating topology in this way will not layout the view.")
        params_config = {
            "topology_type": "star",
            "star_n": nstar,
            "host_counter": 1,
            "switch_counter": 1,
            "link_counter": 1,
            "ip_prefix": subnet
        }
        net_config = kai.create_template_topo(params_config)
        result = kai.deploy_from_config(net_config)
        print(result)


class KlonetFatTreeTopoTemplate(Tool):
    name = "klonet_fattree_topo_template"
    description = ('''
    Deploy a Fat-Tree network topology template on Klonet. Do not call the
    klonet_deploy_network tool if you use this template.
    
    Args:
        subnet (str): The subnet to deploy the topology.
        npod (int, optional): The number of pods of the Fat-Tree (default is 4).
    
    Returns:
        None
        
    Example:
        >>> klonet_fattree_topo_template("192.168.1.0/24", npod=2)
    ''')

    inputs = ["str", "int"]

    @error_handler
    def __call__(self, subnet: str, npod: int = 4):
        print("[Warning] This operation will overwrite the existing topology.")
        kai.reset_project()
        print("[Warning] Creating topology in this way will not layout the view.")
        params_config = {
            "topology_type": "fattree",
            "fattree_k": npod,
            "host_counter": 1,
            "switch_counter": 1,
            "link_counter": 1,
            "ip_prefix": subnet
        }
        net_config = kai.create_template_topo(params_config)
        result = kai.deploy_from_config(net_config)
        print(result)


class KlonetLinearTopoTemplate(Tool):
    name = "klonet_linear_topo_template"
    description = ('''
    Deploy a linear network topology template on Klonet. Do not call the
    klonet_deploy_network tool if you use this template.
    
    Args:
        subnet (str): The subnet to deploy the topology.
        nswitch (int, optional): The number of switches (default is 3).
        nnodes (int, optional): The number of host nodes on each side of 
            the switches (default is 2).
    
    Returns:
        None
    
    Example:
        >>> klonet_linear_topo_template("192.168.1.0/24", nswitch=3, nnodes=2)
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, subset: str, nswitch: int = 3, nnodes: int = 2):
        print("[Warning] This operation will overwrite the existing topology.")
        kai.reset_project()
        print("[Warning] Creating topology in this way will not layout the view.")
        params_config = {
            "topology_type": "linear",
            "linear_m": nswitch,
            "linear_n": nnodes,
            "host_counter": 1,
            "switch_counter": 1,
            "link_counter": 1,
            "ip_prefix": subset
        }
        net_config = kai.create_template_topo(params_config)
        result = kai.deploy_from_config(net_config)
        print(result)


class KlonetConfigurePublicNetworkTool(Tool):
    name = "klonet_configure_public_network"
    description = ('''
    Use this tool to enable external network access for containers. It connects 
    containers to the Internet via the docker0 bridge, utilizing the eth0 network 
    interface.
    
    Args:
        node_name (str): The name of the node.
        turn_on (bool, optional): True for enabling public network access and False
            for disconnecting this node from the Internet (default is True).
    
    Returns:
        None
    
    Example:
        # To connect h1 to the Internet, use:
        >>> klonet_configure_public_network('h1', turn_on=True)
        # To disconnect h1 from the Internet, use:
        >>> klonet_configure_public_network('h1', turn_on=False)
    ''')

    inputs = ["str", "bool"]

    @error_handler
    def __call__(self, node_name: str, turn_on: bool = True):
        result = kai.config_public_network(node_name, turn_on)
        print(result)


class KlonetCheckPublicNetworkTool(Tool):
    name = "klonet_check_public_network"
    description = ('''
    Check if a given node has access to the Internet.
    
    Args:
        node_name (str): The name of the node.
    
    Returns:
        None
    
    Example:
        # 1 means connected and 0 means disconnected.
        >>> klonet_check_public_network('h1')
    ''')

    inputs = ["str"]

    @error_handler
    def __call__(self, node_name: str):
        result = kai.check_public_network(node_name)
        print(result)


class KlonetFileDownloadTool(Tool):
    name = "klonet_file_download"
    description = ('''
    Download file from Klonet nodes. 
    
    Args:
        node_name (str): The name of the node to download from.
        file_name (str): The path of the file to download.
    
    Returns:
        None
    ''')

    inputs = ["str", "str"]

    @error_handler
    def __call__(self, node_name: str, file_name: str):
        print("KlonetAI does not support file download. Please use "
              "Klonet WebUI to download files manually.")


class KlonetFileUploadTool(Tool):
    name = "klonet_file_upload"
    description = ('''
    Upload a file to a given node.
    
    Args:
        node_name (str): The name of the target container node.
        src_filepath (str): The path to the source file to be uploaded.
        tgt_filepath (str, optional): The path on the target node where 
            the file will be stored (default is '/home').
    
    Returns:
        None
    
    Example:
        # Upload the local main.py to the path /home within the container h1.
        >>> klonet_file_upload("/PathTo/main.py", "h1", "/home")
    ''')

    inputs = ["str", "str", "str"]

    @error_handler
    def __call__(self, src_filepath: str, node_name: str, tgt_filepath: str = "/home"):
        result = kai.upload_file(node_name, src_filepath, tgt_filepath)
        print(result)


class KlonetManageWorkerTool(Tool):
    name = "klonet_manage_worker"
    description = ('''
    Register or unregister a physical host machine.
    
    Args:
        worker_ip (str): The IP address of the physical host machine to be 
            registered or unregistered.
        delete_worker (bool, optional): If True, the specified physical host 
            machine will be unregistered. Otherwise, it will be registered
            (default is False).
    
    Returns:
        None
    
    Example:
        # Register a new physical host machine of IP 10.1.1.16 to Klonet system.
        >>> klonet_manage_worker("10.1.1.16", delete_worker=False)
        # Unregister the physical host machine of IP 10.1.1.16 from Klonet system.
        >>> klonet_manage_worker("10.1.1.16", delete_worker=True)
    ''')

    inputs = ["str", "bool"]

    @error_handler
    def __call__(self, worker_ip: str, delete_worker: bool = False):
        result = kai.manage_worker(worker_ip, delete_worker)
        print(result)


class KlonetCheckHealthTool(Tool):
    name = "klonet_check_health"
    description = ('''
    Check the heartbeat health status of all nodes.
    
    Args:
        None
    
    Returns:
        None
    
    Example:
        >>> klonet_check_health()
    ''')

    @error_handler
    def __call__(self):
        is_broken, broken_nodes = kai.check_health()
        print(f"[{'Broken' if is_broken else 'Health'}]", broken_nodes)

tool/klonet.py
Lizonghang/KlonetAI
from .network import NetworkTutorial
from .mxnet import MXNETTutorial

base = (NetworkTutorial,)
mxnet = (MXNETTutorial,)

tutorial/__init__.py
Lizonghang/KlonetAI
from transformers import Tool


class MXNETTutorial(Tool):
    name = "mxnet_tutorial"
    description = ('''
    MXNet supports distributed training enabling us to leverage multiple machines 
    for faster training. In this document, we describe how it works, how to 
    launch a distributed training job and some environment variables which 
    provide more control.
    
    MXNet uses environment variables to assign roles to different processes 
    and to let different processes find the scheduler. The environment 
    variables are required to be set correctly as follows for the training 
    to start:
    
    - DMLC_ROLE: Specifies the role of the process. This can be server, 
        worker or scheduler. Note that there should only be one scheduler.
    - DMLC_PS_ROOT_URI: Specifies the IP of the scheduler.
    - DMLC_PS_ROOT_PORT: Specifies the port that the scheduler listens to.
    - DMLC_NUM_SERVER: Specifies how many server nodes are in the cluster.
    - DMLC_NUM_WORKER: Specifies how many worker nodes are in the cluster.
    
    Below is an example to start 1 scheduler, 1 server, and 2 workers locally:
    # Start a scheduler.
    DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 
    DMLC_NUM_SERVER=1 DMLC_NUM_WORKER=2 python main.py

    # Start a server.
    DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 
    DMLC_NUM_SERVER=1 DMLC_NUM_WORKER=2 python main.py
    
    # For two workers, run the following command for each of them.
    DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 
    DMLC_NUM_SERVER=1 DMLC_NUM_WORKER=2 python main.py
    
    Args:
        None
    
    Returns:
        None
    
    Example:
        >>> mxnet_tutorial()
    ''')

    def __call__(self):
        print(self.description)

tutorial/mxnet.py
Lizonghang/KlonetAI
from transformers import Tool


class NetworkTutorial(Tool):
    name = "network_tutorial"
    description = ('''
    [Q] How to measure data transfer rate / bandwidth between any two hosts ?
    [A] Take host nodes h1 and h2 as an example, follow these two steps:
        1. Start the iperf server on host node h2 as a deamon process.
        2. Launch the iperf client on host node h1. Ensure to convert h2 to
        its IP address as name resolution is not available.
    Args:
        None
        
    Returns:
        None

    Examples:
        >>> network_tutorial()
    ''')

    def __call__(self):
        print(self.description)

tutorial/network.py
Lizonghang/KlonetAI