|
| 1 | +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import json |
| 16 | +import os |
| 17 | + |
| 18 | +from v2.nacos import ClientConfig, NacosConfigService |
| 19 | +from v2.nacos.config.model.config_param import ConfigParam |
| 20 | + |
| 21 | +from veadk.agent import Agent |
| 22 | +from veadk.auth.veauth.mse_veauth import get_mse_cridential |
| 23 | +from veadk.consts import DEFAULT_NACOS_GROUP, DEFAULT_NACOS_INSTANCE_NAME |
| 24 | +from veadk.utils.logger import get_logger |
| 25 | + |
| 26 | +logger = get_logger(__name__) |
| 27 | + |
| 28 | + |
| 29 | +class DynamicConfigManager: |
| 30 | + """ |
| 31 | + DynamicConfigManager is responsible for creating and publishing dynamic config to nacos. |
| 32 | + """ |
| 33 | + |
| 34 | + def __init__( |
| 35 | + self, |
| 36 | + agents: list[Agent] | Agent, |
| 37 | + ): |
| 38 | + """ |
| 39 | + Initialize DynamicConfigManager with agents and app_name. |
| 40 | +
|
| 41 | + Args: |
| 42 | + agents (list[Agent] | Agent): The agent(s) to be included in the dynamic config. |
| 43 | + """ |
| 44 | + if isinstance(agents, list): |
| 45 | + self.agents = agents |
| 46 | + else: |
| 47 | + self.agents = [agents] |
| 48 | + |
| 49 | + logger.debug(f"DynamicConfigManager init with {len(self.agents)} agent(s).") |
| 50 | + |
| 51 | + async def create_config( |
| 52 | + self, |
| 53 | + configs: dict = {}, |
| 54 | + instance_name: str = "", |
| 55 | + group_id: str = "", |
| 56 | + ): |
| 57 | + if not instance_name: |
| 58 | + logger.warning( |
| 59 | + f"instance_name is not provided, use default value `{DEFAULT_NACOS_INSTANCE_NAME}`. This may lead to unexpected behavior such as configuration override." |
| 60 | + ) |
| 61 | + instance_name = DEFAULT_NACOS_INSTANCE_NAME |
| 62 | + |
| 63 | + if not group_id: |
| 64 | + logger.warning( |
| 65 | + f"group_id is not provided, use default value `{DEFAULT_NACOS_GROUP}`. This may lead to unexpected behavior such as configuration override." |
| 66 | + ) |
| 67 | + group_id = group_id or DEFAULT_NACOS_GROUP |
| 68 | + |
| 69 | + nacos_endpoint = os.getenv("NACOS_ENDPOINT") |
| 70 | + nacos_port = os.getenv("NACOS_PORT", "8848") |
| 71 | + nacos_username = os.getenv("NACOS_USERNAME", "nacos") |
| 72 | + nacos_password = os.getenv("NACOS_PASSWORD") |
| 73 | + |
| 74 | + if not all([nacos_endpoint, nacos_port, nacos_username, nacos_password]): |
| 75 | + logger.warning( |
| 76 | + "fetch NACOS_ENDPOINT, NACOS_PORT, NACOS_USERNAME, and NACOS_PASSWORD from env failed, try to get by volcengine AK/SK." |
| 77 | + ) |
| 78 | + |
| 79 | + nacos_credentials = get_mse_cridential(instance_name=instance_name) |
| 80 | + nacos_endpoint = nacos_credentials.endpoint |
| 81 | + nacos_port = nacos_credentials.port |
| 82 | + nacos_username = nacos_credentials.username |
| 83 | + nacos_password = nacos_credentials.password |
| 84 | + |
| 85 | + client_config = ClientConfig( |
| 86 | + server_addresses=f"{nacos_endpoint}:{nacos_port}", |
| 87 | + namespace_id="", |
| 88 | + username=nacos_username, |
| 89 | + password=nacos_password, |
| 90 | + ) |
| 91 | + |
| 92 | + config_client = await NacosConfigService.create_config_service( |
| 93 | + client_config=client_config |
| 94 | + ) |
| 95 | + |
| 96 | + if not configs: |
| 97 | + logger.info("user config_dict is empty, use default config instead.") |
| 98 | + configs = { |
| 99 | + "agent": [ |
| 100 | + { |
| 101 | + "id": agent.id, |
| 102 | + "name": agent.name, |
| 103 | + "description": agent.description, |
| 104 | + "model_name": agent.model_name, |
| 105 | + "instruction": agent.instruction, |
| 106 | + } |
| 107 | + for agent in self.agents |
| 108 | + ] |
| 109 | + } |
| 110 | + response = await config_client.publish_config( |
| 111 | + param=ConfigParam( |
| 112 | + data_id="veadk", |
| 113 | + group=group_id, |
| 114 | + type="json", |
| 115 | + content=json.dumps(configs), |
| 116 | + ) |
| 117 | + ) |
| 118 | + assert response, "publish config to nacos failed" |
| 119 | + logger.info("Publish config to nacos success") |
| 120 | + |
| 121 | + await config_client.add_listener( |
| 122 | + data_id="veadk", |
| 123 | + group="VEADK_GROUP", |
| 124 | + listener=self.handle_config_update, |
| 125 | + ) |
| 126 | + logger.info("Add config listener to nacos success") |
| 127 | + |
| 128 | + return config_client |
| 129 | + |
| 130 | + def register_agent(self, agent: list[Agent] | Agent): |
| 131 | + if isinstance(agent, list): |
| 132 | + self.agents.extend(agent) |
| 133 | + else: |
| 134 | + self.agents.append(agent) |
| 135 | + |
| 136 | + def update_agent(self, configs: dict): |
| 137 | + for agent in self.agents: |
| 138 | + for config in configs["agent"]: |
| 139 | + if agent.id == config["id"]: |
| 140 | + logger.info(f"Update agent {agent.id} with config {config}") |
| 141 | + name = config["name"] |
| 142 | + description = config["description"] |
| 143 | + model_name = config["model_name"] |
| 144 | + instruction = config["instruction"] |
| 145 | + |
| 146 | + agent.name = name |
| 147 | + agent.description = description |
| 148 | + if model_name != agent.model_name: |
| 149 | + agent.update_model(model_name=model_name) |
| 150 | + agent.instruction = instruction |
| 151 | + |
| 152 | + async def handle_config_update(self, tenant, data_id, group, content) -> None: |
| 153 | + logger.debug( |
| 154 | + "listen, tenant:{} data_id:{} group:{} content:{}".format( |
| 155 | + tenant, data_id, group, content |
| 156 | + ) |
| 157 | + ) |
| 158 | + content = json.loads(content) |
| 159 | + self.update_agent(content) |
0 commit comments