From 78e712d3ec83f84155d6a4e2364928c00948a0ea Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Thu, 4 Sep 2025 11:02:50 +0800 Subject: [PATCH 1/6] feat(auth): add auto auth module --- veadk/auth/__init__.py | 13 ++ veadk/auth/base_auth.py | 22 +++ veadk/auth/veauth/__init__.py | 13 ++ veadk/auth/veauth/apmplus_veauth.py | 60 ++++++ veadk/auth/veauth/ark_veauth.py | 76 ++++++++ veadk/auth/veauth/base_veauth.py | 84 +++++++++ veadk/auth/veauth/cozeloop_veauth.py | 13 ++ veadk/auth/veauth/prompt_pilot_veauth.py | 60 ++++++ veadk/auth/veauth/vesearch_veauth.py | 60 ++++++ veadk/configs.py | 229 +++++++++++++++++++++++ veadk/consts.py | 7 + veadk/integrations/utils.py | 53 ++++++ veadk/integrations/ve_tls/__init__.py | 13 ++ veadk/integrations/ve_tls/utils.py | 117 ++++++++++++ veadk/integrations/ve_tls/ve_tls.py | 202 ++++++++++++++++++++ veadk/utils/volcengine_sign.py | 5 +- 16 files changed, 1026 insertions(+), 1 deletion(-) create mode 100644 veadk/auth/__init__.py create mode 100644 veadk/auth/base_auth.py create mode 100644 veadk/auth/veauth/__init__.py create mode 100644 veadk/auth/veauth/apmplus_veauth.py create mode 100644 veadk/auth/veauth/ark_veauth.py create mode 100644 veadk/auth/veauth/base_veauth.py create mode 100644 veadk/auth/veauth/cozeloop_veauth.py create mode 100644 veadk/auth/veauth/prompt_pilot_veauth.py create mode 100644 veadk/auth/veauth/vesearch_veauth.py create mode 100644 veadk/configs.py create mode 100644 veadk/integrations/utils.py create mode 100644 veadk/integrations/ve_tls/__init__.py create mode 100644 veadk/integrations/ve_tls/utils.py create mode 100644 veadk/integrations/ve_tls/ve_tls.py diff --git a/veadk/auth/__init__.py b/veadk/auth/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/auth/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/veadk/auth/base_auth.py b/veadk/auth/base_auth.py new file mode 100644 index 00000000..c42d5e02 --- /dev/null +++ b/veadk/auth/base_auth.py @@ -0,0 +1,22 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +class BaseAuth: + def __init__(self) -> None: ... + + def _fetch_token(self) -> str | dict: ... + + @property + def token(self) -> str | dict: ... diff --git a/veadk/auth/veauth/__init__.py b/veadk/auth/veauth/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/auth/veauth/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/veadk/auth/veauth/apmplus_veauth.py b/veadk/auth/veauth/apmplus_veauth.py new file mode 100644 index 00000000..55d119f3 --- /dev/null +++ b/veadk/auth/veauth/apmplus_veauth.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from typing_extensions import override + +from veadk.auth.veauth.base_veauth import BaseVeAuth +from veadk.utils.logger import get_logger +from veadk.utils.volcengine_sign import ve_request + +logger = get_logger(__name__) + + +class APMPlusVeAuth(BaseVeAuth): + def __init__( + self, + access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), + secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + ) -> None: + super().__init__(access_key, secret_key) + + self._token: str = "" + + @override + def _fetch_token(self) -> None: + logger.info("Fetching APMPlus token...") + + res = ve_request( + request_body={}, + action="GetAppKey", + ak=self.access_key, + sk=self.secret_key, + service="apmplus_server", + version="2024-07-30", + region="cn-beijing", + host="open.volcengineapi.com", + ) + try: + self._token = res["data"]["app_key"] + except KeyError: + raise ValueError(f"Failed to get APMPlus token: {res}") + + @property + def token(self) -> str: + if self._token: + return self._token + self._fetch_token() + return self._token diff --git a/veadk/auth/veauth/ark_veauth.py b/veadk/auth/veauth/ark_veauth.py new file mode 100644 index 00000000..b633ac26 --- /dev/null +++ b/veadk/auth/veauth/ark_veauth.py @@ -0,0 +1,76 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from typing_extensions import override + +from veadk.auth.veauth.base_veauth import BaseVeAuth +from veadk.utils.logger import get_logger +from veadk.utils.volcengine_sign import ve_request + +logger = get_logger(__name__) + + +class ARKVeAuth(BaseVeAuth): + def __init__( + self, + access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), + secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + ) -> None: + super().__init__(access_key, secret_key) + + self._token: str = "" + + @override + def _fetch_token(self) -> None: + # list api keys + first_api_key_id = "" + res = ve_request( + request_body={"ProjectName": "default", "Filter": {}}, + action="ListApiKeys", + ak=self.access_key, + sk=self.secret_key, + service="ark", + version="2024-01-01", + region="cn-beijing", + host="open.volcengineapi.com", + ) + try: + first_api_key_id = res["Result"]["Items"][0]["Id"] + except KeyError: + raise ValueError(f"Failed to get ARK api key list: {res}") + + # get raw api key + res = ve_request( + request_body={"Id": first_api_key_id}, + action="GetRawApiKey", + ak=self.access_key, + sk=self.secret_key, + service="ark", + version="2024-01-01", + region="cn-beijing", + host="open.volcengineapi.com", + ) + try: + self._token = res["Result"]["ApiKey"] + except KeyError: + raise ValueError(f"Failed to get ARK api key: {res}") + + @property + def token(self) -> str: + if self._token: + return self._token + self._fetch_token() + return self._token diff --git a/veadk/auth/veauth/base_veauth.py b/veadk/auth/veauth/base_veauth.py new file mode 100644 index 00000000..f11e8f2e --- /dev/null +++ b/veadk/auth/veauth/base_veauth.py @@ -0,0 +1,84 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from abc import ABC, abstractmethod +from typing import Type + +from veadk.auth.base_auth import BaseAuth + + +class BaseVeAuth(ABC, BaseAuth): + volcengine_access_key: str + """Volcengine Access Key""" + + volcengine_secret_key: str + """Volcengine Secret Key""" + + def __init__( + self, + access_key: str | None = None, + secret_key: str | None = None, + ) -> None: + super().__init__() + + final_ak = access_key or os.getenv("VOLCENGINE_ACCESS_KEY") + final_sk = secret_key or os.getenv("VOLCENGINE_SECRET_KEY") + + assert final_ak, "Volcengine access key cannot be empty." + assert final_sk, "Volcengine secret key cannot be empty." + + self.access_key = final_ak + self.secret_key = final_sk + + self._token: str = "" + + @abstractmethod + def _fetch_token(self) -> None: ... + + @property + def token(self) -> str: ... + + +def veauth(auth_token_name: str, auth_cls: Type[BaseVeAuth]): + def decorator(cls: Type): + # api_key -> _api_key + # for cache + private_auth_token_name = f"_{auth_token_name}" + setattr(cls, private_auth_token_name, "") + + # init a auth cls for fetching token + auth_cls_instance = "_auth_cls_instance" + setattr(cls, auth_cls_instance, auth_cls()) + + def getattribute(self, name: str): + if name != auth_token_name: + return object.__getattribute__(self, name) + if name == auth_token_name: + token = object.__getattribute__(self, name) + + if token: + return token + elif not token and not getattr(cls, private_auth_token_name): + token = getattr(cls, auth_cls_instance).token + setattr(cls, private_auth_token_name, token) + return token + elif not token and getattr(cls, private_auth_token_name): + return getattr(cls, private_auth_token_name) + return token + + setattr(cls, "__getattribute__", getattribute) + return cls + + return decorator diff --git a/veadk/auth/veauth/cozeloop_veauth.py b/veadk/auth/veauth/cozeloop_veauth.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/auth/veauth/cozeloop_veauth.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/veadk/auth/veauth/prompt_pilot_veauth.py b/veadk/auth/veauth/prompt_pilot_veauth.py new file mode 100644 index 00000000..1d4ab909 --- /dev/null +++ b/veadk/auth/veauth/prompt_pilot_veauth.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from typing_extensions import override + +from veadk.auth.veauth.base_veauth import BaseVeAuth +from veadk.utils.logger import get_logger +from veadk.utils.volcengine_sign import ve_request + +logger = get_logger(__name__) + + +class PromptPilotVeAuth(BaseVeAuth): + def __init__( + self, + access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), + secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + ) -> None: + super().__init__(access_key, secret_key) + + self._token: str | dict = "" + + @override + def _fetch_token(self) -> None: + logger.info("Fetching Prompt Pilot token...") + + res = ve_request( + request_body={}, + action="GetOrCreatePromptPilotAPIKeys", + ak=self.access_key, + sk=self.secret_key, + service="ark", + version="2024-01-01", + region="cn-beijing", + host="open.volcengineapi.com", + ) + try: + self._token = res["Result"]["APIKeys"][0]["APIKey"] + except KeyError: + raise ValueError(f"Failed to get Prompt Pilot token: {res}") + + @property + def token(self) -> str | dict: + if self._token: + return self._token + self._fetch_token() + return self._token diff --git a/veadk/auth/veauth/vesearch_veauth.py b/veadk/auth/veauth/vesearch_veauth.py new file mode 100644 index 00000000..920bd001 --- /dev/null +++ b/veadk/auth/veauth/vesearch_veauth.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from typing_extensions import override + +from veadk.auth.veauth.base_veauth import BaseVeAuth +from veadk.utils.logger import get_logger +from veadk.utils.volcengine_sign import ve_request + +logger = get_logger(__name__) + + +class VesearchVeAuth(BaseVeAuth): + def __init__( + self, + access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), + secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + ) -> None: + super().__init__(access_key, secret_key) + + self._token: str = "" + + @override + def _fetch_token(self): + logger.info("Fetching Vesearch token from Volcengine Top Gateway.") + + res = ve_request( + request_body={"biz_scene": "search_agent", "page": 1, "rows": 10}, + action="ListAPIKeys", + ak=self.access_key, + sk=self.secret_key, + service="content_customization", + version="2025-01-01", + region="cn-beijing", + host="open.volcengineapi.com", + ) + try: + self._token = res["Result"]["api_key_vos"][0]["api_key"] + except KeyError: + raise ValueError(f"Failed to get VeSearch token: {res}") + + @property + def token(self) -> str: + if self._token: + return self._token + self._fetch_token() + return self._token diff --git a/veadk/configs.py b/veadk/configs.py new file mode 100644 index 00000000..8c1ab679 --- /dev/null +++ b/veadk/configs.py @@ -0,0 +1,229 @@ +import os + +from dotenv import find_dotenv +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict +from yaml import safe_load + +from veadk.auth.veauth.apmplus_veauth import APMPlusVeAuth +from veadk.auth.veauth.ark_veauth import ARKVeAuth +from veadk.auth.veauth.base_veauth import veauth +from veadk.auth.veauth.prompt_pilot_veauth import PromptPilotVeAuth +from veadk.auth.veauth.vesearch_veauth import VesearchVeAuth +from veadk.consts import ( + DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, + DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, + DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, + DEFAULT_MODEL_AGENT_API_BASE, + DEFAULT_MODEL_AGENT_NAME, + DEFAULT_MODEL_AGENT_PROVIDER, + DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, + DEFAULT_TLS_OTEL_EXPORTER_REGION, + DEFAULT_TOS_BUCKET_NAME, +) +from veadk.integrations.utils import vesource +from veadk.integrations.ve_tls.ve_tls import VeTLS +from veadk.utils.misc import flatten_dict + +veadk_environments: dict = {} + + +def set_envs() -> dict: + config_yaml_path = find_dotenv(filename="config.yaml", usecwd=True) + + with open(config_yaml_path, "r", encoding="utf-8") as yaml_file: + config_dict = safe_load(yaml_file) + + flatten_config_dict = flatten_dict(config_dict) + + for k, v in flatten_config_dict.items(): + global veadk_environments + + k = k.upper() + if k in os.environ: + veadk_environments[k] = os.environ[k] + continue + veadk_environments[k] = str(v) + os.environ[k] = str(v) + + return config_dict + + +config_dict = set_envs() + + +@veauth("api_key", ARKVeAuth) +class ModelConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="MODEL_AGENT_") + + name: str = DEFAULT_MODEL_AGENT_NAME + """Model name for agent reasoning.""" + + provider: str = DEFAULT_MODEL_AGENT_PROVIDER + """Model provider for LiteLLM initialization.""" + + api_base: str = DEFAULT_MODEL_AGENT_API_BASE + """The api base of the model for agent reasoning.""" + + api_key: str = "" + """The api key of the model for agent reasoning.""" + + +@veauth("api_key", PromptPilotVeAuth) +class PromptPilotConfig(BaseModel): + api_key: str = "" + + +@veauth("api_key", VesearchVeAuth) +class VeSearchConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="TOOL_VESEARCH_") + + endpoint: int | str = "" + + api_key: str = "" + + +class BuiltinToolsConfig(BaseModel): + vesearch: VeSearchConfig = Field(default_factory=VeSearchConfig) + + +@veauth("otel_exporter_api_key", APMPlusVeAuth) +class APMPlusConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_ENDPOINT", + ) + + otel_exporter_service_name: str = Field( + default=DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, + alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_SERVICE_NAME", + ) + + otel_exporter_api_key: str = Field( + default="", alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY" + ) + + +class CozeloopConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT", + ) + + otel_exporter_api_key: str = Field( + default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY" + ) + + otel_exporter_space_id: str = Field( + default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME" + ) + + +@vesource( + "otel_exporter_topic_id", + lambda: VeTLS().get_trace_topic_id(), +) +class TLSConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT", + ) + + otel_exporter_region: str = Field( + default=DEFAULT_TLS_OTEL_EXPORTER_REGION, + alias="OBSERVABILITY_OPENTELEMETRY_TLS_REGION", + ) + + otel_exporter_topic_id: str = Field( + default="", + alias="OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME", + ) + + +class VikingKnowledgebaseConfig(BaseModel): + project: str = "default" + """User project in Volcengine console web.""" + + region: str = "cn-beijing" + + +class TOSConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_") + + endpoint: str = "tos-cn-beijing.volces.com" + + region: str = "cn-beijing" + + bucket: str = DEFAULT_TOS_BUCKET_NAME + + +class OpensearchConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_OPENSEARCH_") + + host: str = "" + + port: int = 9200 + + username: str = "" + + password: str = "" + + +class MysqlConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_MYSQL_") + + host: str = "" + + user: str = "" + + password: str = "" + + database: str = "" + + charset: str = "utf8" + + +class RedisConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_REDIS_") + + host: str = "" + + port: int = 6379 + + password: str = "" + + db: int = 0 + + +class PrometheusConfig(BaseSettings): + pushgateway_url: str = "" + + pushgateway_username: str = "" + + pushgateway_password: str = "" + + +class VeADKConfig(BaseModel): + model_agent: ModelConfig = Field(default_factory=ModelConfig) + """The model config for agent reasoning.""" + + tool: BuiltinToolsConfig = Field(default_factory=BuiltinToolsConfig) + """Builtin tools config""" + + apmplus_config: APMPlusConfig = Field(default_factory=APMPlusConfig) + cozeloop_config: CozeloopConfig = Field(default_factory=CozeloopConfig) + tls_config: TLSConfig = Field(default_factory=TLSConfig) + + prometheus_config: PrometheusConfig = Field(default_factory=PrometheusConfig) + + tos: TOSConfig = Field(default_factory=TOSConfig) + + opensearch: OpensearchConfig = Field(default_factory=OpensearchConfig) + mysql: MysqlConfig = Field(default_factory=MysqlConfig) + redis: RedisConfig = Field(default_factory=RedisConfig) + viking_knowledgebase: VikingKnowledgebaseConfig = Field( + default_factory=VikingKnowledgebaseConfig + ) + + +settings = VeADKConfig() diff --git a/veadk/consts.py b/veadk/consts.py index b7be3ec2..509374f0 100644 --- a/veadk/consts.py +++ b/veadk/consts.py @@ -49,3 +49,10 @@ DEFAULT_CR_INSTANCE_NAME = "veadk-user-instance" DEFAULT_CR_NAMESPACE_NAME = "veadk-user-namespace" DEFAULT_CR_REPO_NAME = "veadk-user-repo" + +DEFAULT_TLS_LOG_PROJECT_NAME = "veadk-logs" +DEFAULT_TLS_TRACING_INSTANCE_NAME = "veadk-tracing" + +DEFAULT_TOS_BUCKET_NAME = "veadk-default-bucket" + +DEFAULT_COZELOOP_SPACE_NAME = "veadk-space" diff --git a/veadk/integrations/utils.py b/veadk/integrations/utils.py new file mode 100644 index 00000000..470e026b --- /dev/null +++ b/veadk/integrations/utils.py @@ -0,0 +1,53 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable, Type + + +def vesource(source_name: str, source_func: Callable): + """Automated resource management for Volcengine. + + Args: + source_name (str): The name of the source. + source_func (Callable): The function to call for the source. + + Returns: + str: The result of the source function. + """ + + def decorator(cls: Type): + # record cache source_name -> _source_name + private_source_name = f"_{source_name}" + setattr(cls, private_source_name, "") + + def getattribute(self, name: str): + if name != source_name: + return object.__getattribute__(self, name) + if name == source_name: + source = object.__getattribute__(self, name) + + if source: + return source + elif not source and not getattr(cls, private_source_name): + source = source_func() + setattr(cls, private_source_name, source) + return source + elif not source and getattr(cls, private_source_name): + return getattr(cls, private_source_name) + return source + + setattr(cls, "__getattribute__", getattribute) + return cls + + return decorator diff --git a/veadk/integrations/ve_tls/__init__.py b/veadk/integrations/ve_tls/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/integrations/ve_tls/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/veadk/integrations/ve_tls/utils.py b/veadk/integrations/ve_tls/utils.py new file mode 100644 index 00000000..33009400 --- /dev/null +++ b/veadk/integrations/ve_tls/utils.py @@ -0,0 +1,117 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib +import json +from typing import Any, Literal + +import httpx +from volcengine.ApiInfo import ApiInfo +from volcengine.auth.SignerV4 import SignerV4 +from volcengine.tls.const import ( + APPLICATION_JSON, + CONTENT_MD5, + CONTENT_TYPE, + DATA, + WEB_TRACKS, +) +from volcengine.tls.TLSService import TLSService + +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + +HEADER_API_VERSION = "x-tls-apiversion" +API_VERSION_V_0_3_0 = "0.3.0" + +API_INFO = { + "CreateProject": ApiInfo("POST", "/CreateProject", {}, {}, {}), + "DescribeProjects": ApiInfo("GET", "/DescribeProjects", {}, {}, {}), + "CreateTraceInstance": ApiInfo("POST", "/CreateTraceInstance", {}, {}, {}), + "DescribeTraceInstances": ApiInfo("GET", "/DescribeTraceInstances", {}, {}, {}), + "DescribeTraceInstance": ApiInfo("GET", "/DescribeTraceInstance", {}, {}, {}), +} + + +def __prepare_request( + client: TLSService, + api: str, + params: dict | None = None, + body: Any | None = None, + request_headers: dict | None = None, +): + if params is None: + params = {} + if body is None: + body = {} + + request = client.prepare_request(API_INFO[api], params) + + if request_headers is None: + request_headers = {CONTENT_TYPE: APPLICATION_JSON} + request.headers.update(request_headers) + + if "json" in request.headers[CONTENT_TYPE] and api != WEB_TRACKS: + request.body = json.dumps(body) + else: + request.body = body[DATA] + + if len(request.body) != 0: + if isinstance(request.body, str): + request.headers[CONTENT_MD5] = hashlib.md5( + request.body.encode("utf-8") + ).hexdigest() + else: + request.headers[CONTENT_MD5] = hashlib.md5(request.body).hexdigest() + + SignerV4.sign(request, client.service_info.credentials) + + return request + + +async def ve_tls_request( + client: TLSService, + api: str, + params: dict | None = None, + body: dict | None = None, + request_headers: dict | None = None, + method: Literal["POST", "GET", "PUT", "DELETE"] = "POST", +): + """Customize a standard HTTP request to the Volcengine TLS API""" + + if request_headers is None: + request_headers = {HEADER_API_VERSION: API_VERSION_V_0_3_0} + elif HEADER_API_VERSION not in request_headers: + request_headers[HEADER_API_VERSION] = API_VERSION_V_0_3_0 + if CONTENT_TYPE not in request_headers: + request_headers[CONTENT_TYPE] = APPLICATION_JSON + request = __prepare_request(client, api, params, body, request_headers) + + url = request.build() + + async with httpx.AsyncClient() as session: + while True: + try: + response = await session.request( + method=method, + url=url, + headers=request.headers, + data=request.body, # type: ignore + timeout=60, + ) + return response.json() + except Exception as e: + logger.error( + f"Error occurred while making {method} request to {url}: {e}. Response: {response}" + ) diff --git a/veadk/integrations/ve_tls/ve_tls.py b/veadk/integrations/ve_tls/ve_tls.py new file mode 100644 index 00000000..9871c85d --- /dev/null +++ b/veadk/integrations/ve_tls/ve_tls.py @@ -0,0 +1,202 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os + +from volcengine.tls.TLSService import TLSService + +from veadk.consts import DEFAULT_TLS_LOG_PROJECT_NAME, DEFAULT_TLS_TRACING_INSTANCE_NAME +from veadk.integrations.ve_tls.utils import ve_tls_request +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +class VeTLS: + def __init__( + self, + access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), + secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + region: str = "cn-beijing", + ): + self.access_key = access_key + self.secret_key = secret_key + self.region = region + + self._client = TLSService( + endpoint=f"https://tls-{self.region}.volces.com", + access_key_id=self.access_key, + access_key_secret=self.secret_key, + region=self.region, + ) + + def get_project_id_by_name(self, project_name: str) -> str: + """Get the ID of a log project by its name. + + Args: + project_name (str): The name of the log project. + + Returns: + str: The ID of the log project, or None if not found. + """ + logger.info(f"Getting ID for log project '{project_name}' in TLS...") + + request_body = { + "ProjectName": project_name, + "IsFullName": True, + } + + try: + res = None + res = asyncio.run( + ve_tls_request( + client=self._client, + api="DescribeProjects", + body=request_body, + method="GET", + ) + ) + projects = res["Projects"] + for project in projects: + if project["ProjectName"] == project_name: + return project["ProjectId"] + return "" + except KeyError: + raise ValueError(f"Failed to get log project ID: {res}") + + def create_log_project(self, project_name: str) -> str: + """Create a log project in TLS. + + Args: + project_name (str): The name of the log project to create. + + Returns: + str: The ID of the created log project. + """ + logger.info(f"Creating log project '{project_name}' in TLS...") + + request_body = { + "ProjectName": project_name, + "Region": self.region, + "Description": "Created by Volcengine Agent Development Kit (VeADK)", + "Tags": [{"Key": "provider", "Value": "VeADK"}], + } + try: + res = asyncio.run( + ve_tls_request( + client=self._client, api="CreateProject", body=request_body + ) + ) + + if res["ErrorCode"] == "ProjectAlreadyExist": + logger.debug( + f"Log project '{project_name}' already exists. Check its ID." + ) + return self.get_project_id_by_name(project_name) + + return res["ProjectId"] + except KeyError: + raise ValueError(f"Failed to create log project: {res}") + + def get_trace_instance_by_name(self, log_project_id: str, trace_instance_name: str): + logger.info(f"Getting trace instance '{trace_instance_name}' in TLS...") + + request_body = { + "PageSize": 100, + "ProjectId": log_project_id, + "TraceInstanceName": trace_instance_name, + } + try: + res = asyncio.run( + ve_tls_request( + client=self._client, + api="DescribeTraceInstances", + body=request_body, + method="GET", + ) + ) + + for instance in res["TraceInstances"]: + if instance["TraceInstanceName"] == trace_instance_name: + return instance + except KeyError: + raise ValueError(f"Failed to create log project: {res}") + + def create_tracing_instance(self, log_project_id: str, trace_instance_name: str): + """Create a tracing instance in TLS. + + Args: + instance_name (str): The name of the tracing instance to create. + + Returns: + dict: The tracing instance. + """ + logger.info(f"Creating tracing instance '{trace_instance_name}' in TLS...") + + request_body = { + "ProjectId": log_project_id, + "TraceInstanceName": trace_instance_name, + "Description": "Created by Volcengine Agent Development Kit (VeADK)", + } + + try: + res = None + res = asyncio.run( + ve_tls_request( + client=self._client, + api="CreateTraceInstance", + body=request_body, + ) + ) + + if res["ErrorCode"] == "TopicAlreadyExist": + logger.debug( + f"Log project '{trace_instance_name}' already exists. Check its ID." + ) + return self.get_trace_instance_by_name( + log_project_id, trace_instance_name + ) + + # after creation, get the trace instance details + res = asyncio.run( + ve_tls_request( + client=self._client, + api="DescribeTraceInstance", + body={"TraceInstanceID": res["TraceInstanceID"]}, + method="GET", + ) + ) + + return res + except KeyError: + raise ValueError(f"Failed to create tracing instance: {res}") + + def get_trace_topic_id(self): + """Get the trace topic ID under VeADK default names. + + This method is a tool function just designed for `veadk/config.py`. + + Returns: + str: The trace topic ID. + """ + logger.info("Getting trace topic ID for tracing instance in TLS...") + + log_project_id = self.create_log_project(DEFAULT_TLS_LOG_PROJECT_NAME) + + instance = self.create_tracing_instance( + log_project_id, DEFAULT_TLS_TRACING_INSTANCE_NAME + ) + + return instance["TraceTopicId"] diff --git a/veadk/utils/volcengine_sign.py b/veadk/utils/volcengine_sign.py index 9b82aaed..34d7089d 100644 --- a/veadk/utils/volcengine_sign.py +++ b/veadk/utils/volcengine_sign.py @@ -153,7 +153,10 @@ def request(method, date, query, header, ak, sk, action, body): params=request_param["query"], data=request_param["body"], ) - return r.json() + try: + return r.json() + except Exception: + raise ValueError(f"Error occurred. Bad response: {r}") def ve_request( From 83060deb6ea1f028505ec575c68e7d09f48d9923 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 10 Sep 2025 10:24:21 +0800 Subject: [PATCH 2/6] update auto auth --- tests/test_agent.py | 11 - tests/test_tos.py | 169 ------------- veadk/agent.py | 23 +- veadk/auth/veauth/apmplus_veauth.py | 7 +- veadk/auth/veauth/ark_veauth.py | 1 + veadk/auth/veauth/base_veauth.py | 34 --- veadk/auth/veauth/prompt_pilot_veauth.py | 4 +- veadk/cli/cli_deploy.py | 1 + veadk/cli/cli_prompt.py | 4 +- veadk/config.py | 130 ++++------ veadk/configs.py | 229 ------------------ veadk/configs/__init__.py | 13 + veadk/configs/database_configs.py | 90 +++++++ veadk/configs/model_configs.py | 42 ++++ veadk/configs/tool_configs.py | 42 ++++ veadk/configs/tracing_configs.py | 83 +++++++ veadk/consts.py | 5 +- veadk/integrations/utils.py | 53 ---- .../ve_prompt_pilot/ve_prompt_pilot.py | 9 +- veadk/integrations/ve_tls/ve_tls.py | 18 +- veadk/integrations/ve_tos/ve_tos.py | 128 +++++----- veadk/runner.py | 3 +- veadk/tools/builtin_tools/vesearch.py | 4 +- veadk/tools/builtin_tools/video_generate.py | 47 ++-- .../telemetry/exporters/apmplus_exporter.py | 15 +- .../telemetry/exporters/cozeloop_exporter.py | 13 +- .../telemetry/exporters/tls_exporter.py | 14 +- veadk/utils/logger.py | 2 +- veadk/utils/misc.py | 48 ++++ veadk/utils/volcengine_sign.py | 3 +- 30 files changed, 507 insertions(+), 738 deletions(-) delete mode 100644 tests/test_tos.py delete mode 100644 veadk/configs.py create mode 100644 veadk/configs/__init__.py create mode 100644 veadk/configs/database_configs.py create mode 100644 veadk/configs/model_configs.py create mode 100644 veadk/configs/tool_configs.py create mode 100644 veadk/configs/tracing_configs.py delete mode 100644 veadk/integrations/utils.py diff --git a/tests/test_agent.py b/tests/test_agent.py index 02bb8b6a..45c23347 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -199,17 +199,6 @@ def test_agent_with_tracers(): assert tracer2 in agent.tracers -@patch.dict( - "os.environ", - {"MODEL_AGENT_NAME": "env_model_name", "MODEL_AGENT_API_KEY": "mock_api_key"}, - clear=True, -) -def test_agent_environment_variables(): - agent = Agent() - print(agent) - assert agent.model_name == "env_model_name" - - @patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"}) def test_agent_custom_name_and_description(): custom_name = "CustomAgent" diff --git a/tests/test_tos.py b/tests/test_tos.py deleted file mode 100644 index 6490d89d..00000000 --- a/tests/test_tos.py +++ /dev/null @@ -1,169 +0,0 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import pytest -from unittest import mock - -# Check if tos module is available -import importlib - -TOS_AVAILABLE = False -try: - importlib.import_module("veadk.integrations.ve_tos.ve_tos") - TOS_AVAILABLE = True -except ImportError: - pass - -# Skip tests that require tos module if it's not available -require_tos = pytest.mark.skipif(not TOS_AVAILABLE, reason="tos module not available") - -# 使用 pytest-asyncio -pytest_plugins = ("pytest_asyncio",) - - -@pytest.fixture -@require_tos -def mock_client(monkeypatch): - import veadk.integrations.ve_tos.ve_tos as tos_mod - - fake_client = mock.Mock() - - monkeypatch.setenv("DATABASE_TOS_REGION", "test-region") - monkeypatch.setenv("VOLCENGINE_ACCESS_KEY", "test-access-key") - monkeypatch.setenv("VOLCENGINE_SECRET_KEY", "test-secret-key") - monkeypatch.setenv("DATABASE_TOS_BUCKET", "test-bucket") - - monkeypatch.setattr(tos_mod.tos, "TosClientV2", lambda *a, **k: fake_client) - - class FakeExceptions: - class TosServerError(Exception): - def __init__( - self, - msg: str, - code: int = 0, - host_id: str = "", - resource: str = "", - request_id: str = "", - header=None, - ): - super().__init__(msg) - self.status_code = code - - monkeypatch.setattr(tos_mod.tos, "exceptions", FakeExceptions) - monkeypatch.setattr( - tos_mod.tos, - "StorageClassType", - type("S", (), {"Storage_Class_Standard": "STANDARD"}), - ) - monkeypatch.setattr( - tos_mod.tos, - "ACLType", - type("A", (), {"ACL_Private": "private", "ACL_Public_Read": "public-read"}), - ) - - return fake_client - - -@pytest.fixture -@require_tos -def tos_client(mock_client): - import veadk.integrations.ve_tos.ve_tos as tos_mod - - return tos_mod.VeTOS() - - -@require_tos -def test_create_bucket_exists(tos_client, mock_client): - mock_client.head_bucket.return_value = None # head_bucket 正常返回表示存在 - result = tos_client.create_bucket() - assert result is True - mock_client.create_bucket.assert_not_called() - - -@require_tos -def test_create_bucket_not_exists(tos_client, mock_client): - import veadk.integrations.ve_tos.ve_tos as tos_mod - - exc = tos_mod.tos.exceptions.TosServerError(msg="not found", code=404) - mock_client.head_bucket.side_effect = exc - - result = tos_client.create_bucket() - assert result is True - mock_client.create_bucket.assert_called_once() - - -@require_tos -@pytest.mark.asyncio -async def test_upload_bytes_success(tos_client, mock_client): - mock_client.head_bucket.return_value = True - data = b"hello world" - - result = await tos_client.upload("obj-key", data) - assert result is None - mock_client.put_object.assert_called_once() - mock_client.close.assert_called_once() - - -@require_tos -@pytest.mark.asyncio -async def test_upload_file_success(tmp_path, tos_client, mock_client): - mock_client.head_bucket.return_value = True - file_path = tmp_path / "file.txt" - file_path.write_text("hello file") - - result = await tos_client.upload("obj-key", str(file_path)) - assert result is None - mock_client.put_object_from_file.assert_called_once() - mock_client.close.assert_called_once() - - -@require_tos -def test_download_success(tmp_path, tos_client, mock_client): - save_path = tmp_path / "out.txt" - mock_client.get_object.return_value = [b"abc", b"def"] - - result = tos_client.download("obj-key", str(save_path)) - assert result is True - assert save_path.read_bytes() == b"abcdef" - - -@require_tos -def test_download_fail(tos_client, mock_client): - mock_client.get_object.side_effect = Exception("boom") - result = tos_client.download("obj-key", "somewhere.txt") - assert result is False - - -@require_tos -@pytest.mark.skipif(TOS_AVAILABLE, reason="tos module is available") -def test_tos_import_error(): - """Test VeTOS behavior when tos module is not installed""" - # Remove tos from sys.modules to simulate it's not installed - import sys - - original_tos = sys.modules.get("tos") - if "tos" in sys.modules: - del sys.modules["tos"] - - try: - # Try to import ve_tos module, which should raise ImportError - with pytest.raises(ImportError) as exc_info: - pass - - # Check that the error message contains installation instructions - assert "pip install tos" in str(exc_info.value) - finally: - # Restore original state - if original_tos is not None: - sys.modules["tos"] = original_tos diff --git a/veadk/agent.py b/veadk/agent.py index fc0e6c63..0af9f586 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -26,12 +26,9 @@ from pydantic import ConfigDict, Field from typing_extensions import Any -from veadk.config import getenv +from veadk.config import settings from veadk.consts import ( DEFAULT_AGENT_NAME, - DEFAULT_MODEL_AGENT_API_BASE, - DEFAULT_MODEL_AGENT_NAME, - DEFAULT_MODEL_AGENT_PROVIDER, DEFAULT_MODEL_EXTRA_CONFIG, ) from veadk.evaluation import EvalSetRecorder @@ -63,26 +60,16 @@ class Agent(LlmAgent): instruction: str = DEFAULT_INSTRUCTION """The instruction for the agent, such as principles of function calling.""" - model_name: str = Field( - default_factory=lambda: getenv("MODEL_AGENT_NAME", DEFAULT_MODEL_AGENT_NAME) - ) + model_name: str = Field(default_factory=lambda: settings.model.name) """The name of the model for agent running.""" - model_provider: str = Field( - default_factory=lambda: getenv( - "MODEL_AGENT_PROVIDER", DEFAULT_MODEL_AGENT_PROVIDER - ) - ) + model_provider: str = Field(default_factory=lambda: settings.model.provider) """The provider of the model for agent running.""" - model_api_base: str = Field( - default_factory=lambda: getenv( - "MODEL_AGENT_API_BASE", DEFAULT_MODEL_AGENT_API_BASE - ) - ) + model_api_base: str = Field(default_factory=lambda: settings.model.api_base) """The api base of the model for agent running.""" - model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY")) + model_api_key: str = Field(default_factory=lambda: settings.model.api_key) """The api key of the model for agent running.""" model_extra_config: dict = Field(default_factory=dict) diff --git a/veadk/auth/veauth/apmplus_veauth.py b/veadk/auth/veauth/apmplus_veauth.py index 55d119f3..9d15dc9e 100644 --- a/veadk/auth/veauth/apmplus_veauth.py +++ b/veadk/auth/veauth/apmplus_veauth.py @@ -28,9 +28,12 @@ def __init__( self, access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + region: str = "cn-beijing", ) -> None: super().__init__(access_key, secret_key) + self.region = region + self._token: str = "" @override @@ -44,8 +47,10 @@ def _fetch_token(self) -> None: sk=self.secret_key, service="apmplus_server", version="2024-07-30", - region="cn-beijing", + region=self.region, host="open.volcengineapi.com", + # APMPlus frontend required + header={"X-Apmplus-Region": self.region.replace("-", "_")}, ) try: self._token = res["data"]["app_key"] diff --git a/veadk/auth/veauth/ark_veauth.py b/veadk/auth/veauth/ark_veauth.py index b633ac26..ceae30be 100644 --- a/veadk/auth/veauth/ark_veauth.py +++ b/veadk/auth/veauth/ark_veauth.py @@ -35,6 +35,7 @@ def __init__( @override def _fetch_token(self) -> None: + logger.info("Fetching ARK token...") # list api keys first_api_key_id = "" res = ve_request( diff --git a/veadk/auth/veauth/base_veauth.py b/veadk/auth/veauth/base_veauth.py index f11e8f2e..4a66d69d 100644 --- a/veadk/auth/veauth/base_veauth.py +++ b/veadk/auth/veauth/base_veauth.py @@ -14,7 +14,6 @@ import os from abc import ABC, abstractmethod -from typing import Type from veadk.auth.base_auth import BaseAuth @@ -49,36 +48,3 @@ def _fetch_token(self) -> None: ... @property def token(self) -> str: ... - - -def veauth(auth_token_name: str, auth_cls: Type[BaseVeAuth]): - def decorator(cls: Type): - # api_key -> _api_key - # for cache - private_auth_token_name = f"_{auth_token_name}" - setattr(cls, private_auth_token_name, "") - - # init a auth cls for fetching token - auth_cls_instance = "_auth_cls_instance" - setattr(cls, auth_cls_instance, auth_cls()) - - def getattribute(self, name: str): - if name != auth_token_name: - return object.__getattribute__(self, name) - if name == auth_token_name: - token = object.__getattribute__(self, name) - - if token: - return token - elif not token and not getattr(cls, private_auth_token_name): - token = getattr(cls, auth_cls_instance).token - setattr(cls, private_auth_token_name, token) - return token - elif not token and getattr(cls, private_auth_token_name): - return getattr(cls, private_auth_token_name) - return token - - setattr(cls, "__getattribute__", getattribute) - return cls - - return decorator diff --git a/veadk/auth/veauth/prompt_pilot_veauth.py b/veadk/auth/veauth/prompt_pilot_veauth.py index 1d4ab909..c2aeab58 100644 --- a/veadk/auth/veauth/prompt_pilot_veauth.py +++ b/veadk/auth/veauth/prompt_pilot_veauth.py @@ -31,7 +31,7 @@ def __init__( ) -> None: super().__init__(access_key, secret_key) - self._token: str | dict = "" + self._token: str = "" @override def _fetch_token(self) -> None: @@ -53,7 +53,7 @@ def _fetch_token(self) -> None: raise ValueError(f"Failed to get Prompt Pilot token: {res}") @property - def token(self) -> str | dict: + def token(self) -> str: if self._token: return self._token self._fetch_token() diff --git a/veadk/cli/cli_deploy.py b/veadk/cli/cli_deploy.py index 56bb53f4..293e1f20 100644 --- a/veadk/cli/cli_deploy.py +++ b/veadk/cli/cli_deploy.py @@ -14,6 +14,7 @@ import click + from veadk.version import VERSION TEMP_PATH = "/tmp" diff --git a/veadk/cli/cli_prompt.py b/veadk/cli/cli_prompt.py index c7df61e3..5f05b6cb 100644 --- a/veadk/cli/cli_prompt.py +++ b/veadk/cli/cli_prompt.py @@ -31,7 +31,7 @@ def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None: from pathlib import Path from veadk.agent import Agent - from veadk.config import getenv + from veadk.config import settings from veadk.integrations.ve_prompt_pilot.ve_prompt_pilot import VePromptPilot from veadk.utils.misc import load_module_from_file @@ -55,7 +55,7 @@ def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None: click.echo(f"Found {len(agents)} agents in {module_abs_path}") if not api_key: - api_key = getenv("PROMPT_PILOT_API_KEY") + api_key = settings.prompt_pilot.api_key ve_prompt_pilot = VePromptPilot(api_key) ve_prompt_pilot.optimize( agents=agents, feedback=feedback, model_name=model_name diff --git a/veadk/config.py b/veadk/config.py index 8aee7c0c..6ceeb076 100644 --- a/veadk/config.py +++ b/veadk/config.py @@ -13,96 +13,45 @@ # limitations under the License. import os -from typing import Any, Dict, List, MutableMapping, Tuple +from typing import Any from dotenv import find_dotenv -from pydantic_settings import ( - BaseSettings, - InitSettingsSource, - PydanticBaseSettingsSource, - SettingsConfigDict, - YamlConfigSettingsSource, +from pydantic import BaseModel, Field + +from veadk.configs.database_configs import ( + MysqlConfig, + OpensearchConfig, + PrometheusConfig, + RedisConfig, + TOSConfig, + VikingKnowledgebaseConfig, ) - -settings = None -veadk_environments = {} - - -def flatten_dict( - d: MutableMapping[str, Any], parent_key: str = "", sep: str = "_" -) -> Dict[str, Any]: - """Flatten a nested dictionary, using a separator in the keys. - Useful for pydantic_v1 models with nested fields -- first use - dct = mdl.model_dump() - to get a nested dictionary, then use this function to flatten it. - """ - items: List[Tuple[str, Any]] = [] - for k, v in d.items(): - new_key = f"{parent_key}{sep}{k}" if parent_key else k - if isinstance(v, MutableMapping): - items.extend(flatten_dict(v, new_key, sep=sep).items()) - else: - items.append((new_key, v)) - return dict(items) - - -class Settings(BaseSettings): - model_config = SettingsConfigDict( - yaml_file=find_dotenv(filename="config.yaml", usecwd=True), extra="allow" +from veadk.configs.model_configs import ModelConfig +from veadk.configs.tool_configs import BuiltinToolConfigs, PromptPilotConfig +from veadk.configs.tracing_configs import APMPlusConfig, CozeloopConfig, TLSConfig +from veadk.utils.misc import set_envs + + +class VeADKConfig(BaseModel): + model: ModelConfig = Field(default_factory=ModelConfig) + """Config for agent reasoning model.""" + + tool: BuiltinToolConfigs = Field(default_factory=BuiltinToolConfigs) + prompt_pilot: PromptPilotConfig = Field(default_factory=PromptPilotConfig) + + apmplus_config: APMPlusConfig = Field(default_factory=APMPlusConfig) + cozeloop_config: CozeloopConfig = Field(default_factory=CozeloopConfig) + tls_config: TLSConfig = Field(default_factory=TLSConfig) + prometheus_config: PrometheusConfig = Field(default_factory=PrometheusConfig) + + tos: TOSConfig = Field(default_factory=TOSConfig) + opensearch: OpensearchConfig = Field(default_factory=OpensearchConfig) + mysql: MysqlConfig = Field(default_factory=MysqlConfig) + redis: RedisConfig = Field(default_factory=RedisConfig) + viking_knowledgebase: VikingKnowledgebaseConfig = Field( + default_factory=VikingKnowledgebaseConfig ) - @classmethod - def settings_customise_sources( - cls, - settings_cls: type[BaseSettings], - init_settings: PydanticBaseSettingsSource, - env_settings: PydanticBaseSettingsSource, - dotenv_settings: PydanticBaseSettingsSource, - file_secret_settings: PydanticBaseSettingsSource, - ) -> tuple[PydanticBaseSettingsSource, ...]: - yaml_source = YamlConfigSettingsSource(settings_cls) - raw_data = yaml_source() - flat_data = flatten_dict(raw_data) - - init_source = InitSettingsSource(settings_cls, flat_data) - return ( - init_source, - env_settings, - dotenv_settings, - file_secret_settings, - ) - - -def prepare_settings(): - path = find_dotenv(filename="config.yaml", usecwd=True) - - if path == "" or path is None or not os.path.exists(path): - # logger.warning( - # "Default and recommanded config file `config.yaml` not found. Please put it in the root directory of your project." - # ) - pass - else: - # logger.info(f"Loading config file from {path}") - global settings - settings = Settings() - - for k, v in settings.model_dump().items(): - global veadk_environments - - k = k.upper() - if k in os.environ: - veadk_environments[k] = os.environ[k] - continue - veadk_environments[k] = str(v) - os.environ[k] = str(v) - - -prepare_settings() - - -def get_envlist(): - return os.environ.keys() - def getenv( env_name: str, default_value: Any = "", allow_false_values: bool = False @@ -129,3 +78,14 @@ def getenv( raise ValueError( f"The environment variable `{env_name}` not exists. Please set this in your environment variable or config.yaml." ) + + +config_yaml_path = find_dotenv(filename="config.yaml", usecwd=True) + +veadk_environments = {} + +if config_yaml_path: + config_dict, _veadk_environments = set_envs(config_yaml_path=config_yaml_path) + veadk_environments.update(_veadk_environments) + +settings = VeADKConfig() diff --git a/veadk/configs.py b/veadk/configs.py deleted file mode 100644 index 8c1ab679..00000000 --- a/veadk/configs.py +++ /dev/null @@ -1,229 +0,0 @@ -import os - -from dotenv import find_dotenv -from pydantic import BaseModel, Field -from pydantic_settings import BaseSettings, SettingsConfigDict -from yaml import safe_load - -from veadk.auth.veauth.apmplus_veauth import APMPlusVeAuth -from veadk.auth.veauth.ark_veauth import ARKVeAuth -from veadk.auth.veauth.base_veauth import veauth -from veadk.auth.veauth.prompt_pilot_veauth import PromptPilotVeAuth -from veadk.auth.veauth.vesearch_veauth import VesearchVeAuth -from veadk.consts import ( - DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, - DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, - DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, - DEFAULT_MODEL_AGENT_API_BASE, - DEFAULT_MODEL_AGENT_NAME, - DEFAULT_MODEL_AGENT_PROVIDER, - DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, - DEFAULT_TLS_OTEL_EXPORTER_REGION, - DEFAULT_TOS_BUCKET_NAME, -) -from veadk.integrations.utils import vesource -from veadk.integrations.ve_tls.ve_tls import VeTLS -from veadk.utils.misc import flatten_dict - -veadk_environments: dict = {} - - -def set_envs() -> dict: - config_yaml_path = find_dotenv(filename="config.yaml", usecwd=True) - - with open(config_yaml_path, "r", encoding="utf-8") as yaml_file: - config_dict = safe_load(yaml_file) - - flatten_config_dict = flatten_dict(config_dict) - - for k, v in flatten_config_dict.items(): - global veadk_environments - - k = k.upper() - if k in os.environ: - veadk_environments[k] = os.environ[k] - continue - veadk_environments[k] = str(v) - os.environ[k] = str(v) - - return config_dict - - -config_dict = set_envs() - - -@veauth("api_key", ARKVeAuth) -class ModelConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="MODEL_AGENT_") - - name: str = DEFAULT_MODEL_AGENT_NAME - """Model name for agent reasoning.""" - - provider: str = DEFAULT_MODEL_AGENT_PROVIDER - """Model provider for LiteLLM initialization.""" - - api_base: str = DEFAULT_MODEL_AGENT_API_BASE - """The api base of the model for agent reasoning.""" - - api_key: str = "" - """The api key of the model for agent reasoning.""" - - -@veauth("api_key", PromptPilotVeAuth) -class PromptPilotConfig(BaseModel): - api_key: str = "" - - -@veauth("api_key", VesearchVeAuth) -class VeSearchConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="TOOL_VESEARCH_") - - endpoint: int | str = "" - - api_key: str = "" - - -class BuiltinToolsConfig(BaseModel): - vesearch: VeSearchConfig = Field(default_factory=VeSearchConfig) - - -@veauth("otel_exporter_api_key", APMPlusVeAuth) -class APMPlusConfig(BaseSettings): - otel_exporter_endpoint: str = Field( - default=DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, - alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_ENDPOINT", - ) - - otel_exporter_service_name: str = Field( - default=DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, - alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_SERVICE_NAME", - ) - - otel_exporter_api_key: str = Field( - default="", alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY" - ) - - -class CozeloopConfig(BaseSettings): - otel_exporter_endpoint: str = Field( - default=DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, - alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT", - ) - - otel_exporter_api_key: str = Field( - default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY" - ) - - otel_exporter_space_id: str = Field( - default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME" - ) - - -@vesource( - "otel_exporter_topic_id", - lambda: VeTLS().get_trace_topic_id(), -) -class TLSConfig(BaseSettings): - otel_exporter_endpoint: str = Field( - default=DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, - alias="OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT", - ) - - otel_exporter_region: str = Field( - default=DEFAULT_TLS_OTEL_EXPORTER_REGION, - alias="OBSERVABILITY_OPENTELEMETRY_TLS_REGION", - ) - - otel_exporter_topic_id: str = Field( - default="", - alias="OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME", - ) - - -class VikingKnowledgebaseConfig(BaseModel): - project: str = "default" - """User project in Volcengine console web.""" - - region: str = "cn-beijing" - - -class TOSConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_") - - endpoint: str = "tos-cn-beijing.volces.com" - - region: str = "cn-beijing" - - bucket: str = DEFAULT_TOS_BUCKET_NAME - - -class OpensearchConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="DATABASE_OPENSEARCH_") - - host: str = "" - - port: int = 9200 - - username: str = "" - - password: str = "" - - -class MysqlConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="DATABASE_MYSQL_") - - host: str = "" - - user: str = "" - - password: str = "" - - database: str = "" - - charset: str = "utf8" - - -class RedisConfig(BaseSettings): - model_config = SettingsConfigDict(env_prefix="DATABASE_REDIS_") - - host: str = "" - - port: int = 6379 - - password: str = "" - - db: int = 0 - - -class PrometheusConfig(BaseSettings): - pushgateway_url: str = "" - - pushgateway_username: str = "" - - pushgateway_password: str = "" - - -class VeADKConfig(BaseModel): - model_agent: ModelConfig = Field(default_factory=ModelConfig) - """The model config for agent reasoning.""" - - tool: BuiltinToolsConfig = Field(default_factory=BuiltinToolsConfig) - """Builtin tools config""" - - apmplus_config: APMPlusConfig = Field(default_factory=APMPlusConfig) - cozeloop_config: CozeloopConfig = Field(default_factory=CozeloopConfig) - tls_config: TLSConfig = Field(default_factory=TLSConfig) - - prometheus_config: PrometheusConfig = Field(default_factory=PrometheusConfig) - - tos: TOSConfig = Field(default_factory=TOSConfig) - - opensearch: OpensearchConfig = Field(default_factory=OpensearchConfig) - mysql: MysqlConfig = Field(default_factory=MysqlConfig) - redis: RedisConfig = Field(default_factory=RedisConfig) - viking_knowledgebase: VikingKnowledgebaseConfig = Field( - default_factory=VikingKnowledgebaseConfig - ) - - -settings = VeADKConfig() diff --git a/veadk/configs/__init__.py b/veadk/configs/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/configs/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/veadk/configs/database_configs.py b/veadk/configs/database_configs.py new file mode 100644 index 00000000..d8879cb5 --- /dev/null +++ b/veadk/configs/database_configs.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from functools import cached_property + +from pydantic import BaseModel +from pydantic_settings import BaseSettings, SettingsConfigDict + +from veadk.consts import DEFAULT_TOS_BUCKET_NAME +from veadk.integrations.ve_tos.ve_tos import VeTOS + + +class OpensearchConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_OPENSEARCH_") + + host: str = "" + + port: int = 9200 + + username: str = "" + + password: str = "" + + +class MysqlConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_MYSQL_") + + host: str = "" + + user: str = "" + + password: str = "" + + database: str = "" + + charset: str = "utf8" + + +class RedisConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_REDIS_") + + host: str = "" + + port: int = 6379 + + password: str = "" + + db: int = 0 + + +class PrometheusConfig(BaseSettings): + pushgateway_url: str = "" + + pushgateway_username: str = "" + + pushgateway_password: str = "" + + +class VikingKnowledgebaseConfig(BaseModel): + project: str = "default" + """User project in Volcengine console web.""" + + region: str = "cn-beijing" + + +class TOSConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_") + + endpoint: str = "tos-cn-beijing.volces.com" + + region: str = "cn-beijing" + + @cached_property + def bucket(self) -> str: + _bucket = os.getenv("DATABASE_TOS_BUCKET") or DEFAULT_TOS_BUCKET_NAME + + VeTOS(bucket_name=_bucket).create_bucket() + return _bucket diff --git a/veadk/configs/model_configs.py b/veadk/configs/model_configs.py new file mode 100644 index 00000000..e0efbdb6 --- /dev/null +++ b/veadk/configs/model_configs.py @@ -0,0 +1,42 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from functools import cached_property + +from pydantic_settings import BaseSettings, SettingsConfigDict + +from veadk.auth.veauth.ark_veauth import ARKVeAuth +from veadk.consts import ( + DEFAULT_MODEL_AGENT_API_BASE, + DEFAULT_MODEL_AGENT_NAME, + DEFAULT_MODEL_AGENT_PROVIDER, +) + + +class ModelConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="MODEL_AGENT_") + + name: str = DEFAULT_MODEL_AGENT_NAME + """Model name for agent reasoning.""" + + provider: str = DEFAULT_MODEL_AGENT_PROVIDER + """Model provider for LiteLLM initialization.""" + + api_base: str = DEFAULT_MODEL_AGENT_API_BASE + """The api base of the model for agent reasoning.""" + + @cached_property + def api_key(self) -> str: + return os.getenv("MODEL_AGENT_API_KEY") or ARKVeAuth().token diff --git a/veadk/configs/tool_configs.py b/veadk/configs/tool_configs.py new file mode 100644 index 00000000..8cca6971 --- /dev/null +++ b/veadk/configs/tool_configs.py @@ -0,0 +1,42 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from functools import cached_property + +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from veadk.auth.veauth.prompt_pilot_veauth import PromptPilotVeAuth +from veadk.auth.veauth.vesearch_veauth import VesearchVeAuth + + +class PromptPilotConfig(BaseModel): + @cached_property + def api_key(self) -> str: + return os.getenv("PROMPT_PILOT_API_KEY") or PromptPilotVeAuth().token + + +class VeSearchConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="TOOL_VESEARCH_") + + endpoint: int | str = "" + + @cached_property + def api_key(self) -> str: + return os.getenv("TOOL_VESEARCH_API_KEY") or VesearchVeAuth().token + + +class BuiltinToolConfigs(BaseModel): + vesearch: VeSearchConfig = Field(default_factory=VeSearchConfig) diff --git a/veadk/configs/tracing_configs.py b/veadk/configs/tracing_configs.py new file mode 100644 index 00000000..a53c5d72 --- /dev/null +++ b/veadk/configs/tracing_configs.py @@ -0,0 +1,83 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from functools import cached_property + +from pydantic import Field +from pydantic_settings import BaseSettings + +from veadk.auth.veauth.apmplus_veauth import APMPlusVeAuth +from veadk.consts import ( + DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, + DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, + DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, + DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, + DEFAULT_TLS_OTEL_EXPORTER_REGION, +) +from veadk.integrations.ve_tls.ve_tls import VeTLS + + +class APMPlusConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_ENDPOINT", + ) + + otel_exporter_service_name: str = Field( + default=DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME, + alias="OBSERVABILITY_OPENTELEMETRY_APMPLUS_SERVICE_NAME", + ) + + @cached_property + def otel_exporter_api_key(self) -> str: + return ( + os.getenv("OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY") + or APMPlusVeAuth().token + ) + + +class CozeloopConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT", + ) + + otel_exporter_api_key: str = Field( + default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY" + ) + + otel_exporter_space_id: str = Field( + default="", alias="OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME" + ) + + +class TLSConfig(BaseSettings): + otel_exporter_endpoint: str = Field( + default=DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT, + alias="OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT", + ) + + otel_exporter_region: str = Field( + default=DEFAULT_TLS_OTEL_EXPORTER_REGION, + alias="OBSERVABILITY_OPENTELEMETRY_TLS_REGION", + ) + + @cached_property + def otel_exporter_topic_id(self) -> str: + _topic_id = ( + os.getenv("OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME") + or VeTLS().get_trace_topic_id() + ) + return _topic_id diff --git a/veadk/consts.py b/veadk/consts.py index 509374f0..3163a52f 100644 --- a/veadk/consts.py +++ b/veadk/consts.py @@ -14,7 +14,7 @@ import time -from veadk.config import getenv +from veadk.utils.misc import getenv from veadk.version import VERSION DEFAULT_AGENT_NAME = "veAgent" @@ -32,6 +32,9 @@ "caching": { "type": getenv("MODEL_AGENT_CACHING", "enabled"), }, + "thinking": { + "type": "disabled" + }, # disable thinking for token saving and shorter TTFT by default "expire_at": int(time.time()) + 3600, # expire after 1 hour }, } diff --git a/veadk/integrations/utils.py b/veadk/integrations/utils.py deleted file mode 100644 index 470e026b..00000000 --- a/veadk/integrations/utils.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Callable, Type - - -def vesource(source_name: str, source_func: Callable): - """Automated resource management for Volcengine. - - Args: - source_name (str): The name of the source. - source_func (Callable): The function to call for the source. - - Returns: - str: The result of the source function. - """ - - def decorator(cls: Type): - # record cache source_name -> _source_name - private_source_name = f"_{source_name}" - setattr(cls, private_source_name, "") - - def getattribute(self, name: str): - if name != source_name: - return object.__getattribute__(self, name) - if name == source_name: - source = object.__getattribute__(self, name) - - if source: - return source - elif not source and not getattr(cls, private_source_name): - source = source_func() - setattr(cls, private_source_name, source) - return source - elif not source and getattr(cls, private_source_name): - return getattr(cls, private_source_name) - return source - - setattr(cls, "__getattribute__", getattribute) - return cls - - return decorator diff --git a/veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py b/veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py index b04ffa31..41e9f0f1 100644 --- a/veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py +++ b/veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py @@ -65,12 +65,15 @@ def optimize( api_key=self.api_key, ): # stream chunks of optimized prompt # Process each chunk as it arrives - optimized_prompt += chunk.data.content + optimized_prompt += chunk.data.content if chunk.data else "" # print(chunk.data.content, end="", flush=True) if chunk.event == "usage": - usage = chunk.data.usage + usage = chunk.data.usage if chunk.data else 0 optimized_prompt = optimized_prompt.replace("\\n", "\n") print(f"Optimized prompt for agent {agent.name}:\n{optimized_prompt}") - logger.info(f"Token usage: {usage['total_tokens']}") + if usage: + logger.info(f"Token usage: {usage['total_tokens']}") + else: + logger.warning("No usage data.") return optimized_prompt diff --git a/veadk/integrations/ve_tls/ve_tls.py b/veadk/integrations/ve_tls/ve_tls.py index 9871c85d..3828eda6 100644 --- a/veadk/integrations/ve_tls/ve_tls.py +++ b/veadk/integrations/ve_tls/ve_tls.py @@ -15,11 +15,10 @@ import asyncio import os -from volcengine.tls.TLSService import TLSService - from veadk.consts import DEFAULT_TLS_LOG_PROJECT_NAME, DEFAULT_TLS_TRACING_INSTANCE_NAME from veadk.integrations.ve_tls.utils import ve_tls_request from veadk.utils.logger import get_logger +from volcengine.tls.TLSService import TLSService logger = get_logger(__name__) @@ -27,12 +26,16 @@ class VeTLS: def __init__( self, - access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""), - secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""), + access_key: str | None = None, + secret_key: str | None = None, region: str = "cn-beijing", ): - self.access_key = access_key - self.secret_key = secret_key + self.access_key = ( + access_key if access_key else os.getenv("VOLCENGINE_ACCESS_KEY", "") + ) + self.secret_key = ( + secret_key if secret_key else os.getenv("VOLCENGINE_SECRET_KEY", "") + ) self.region = region self._client = TLSService( @@ -199,4 +202,7 @@ def get_trace_topic_id(self): log_project_id, DEFAULT_TLS_TRACING_INSTANCE_NAME ) + if not instance: + raise ValueError("None instance") + return instance["TraceTopicId"] diff --git a/veadk/integrations/ve_tos/ve_tos.py b/veadk/integrations/ve_tos/ve_tos.py index 89162d16..a52a952b 100644 --- a/veadk/integrations/ve_tos/ve_tos.py +++ b/veadk/integrations/ve_tos/ve_tos.py @@ -12,75 +12,67 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -from veadk.config import getenv -from veadk.utils.logger import get_logger import asyncio -from typing import Union -from pydantic import BaseModel, Field -from typing import Any -from urllib.parse import urlparse +import os from datetime import datetime +from typing import TYPE_CHECKING, Union +from urllib.parse import urlparse + +from veadk.consts import DEFAULT_TOS_BUCKET_NAME +from veadk.utils.logger import get_logger + +if TYPE_CHECKING: + import tos + # Initialize logger before using it logger = get_logger(__name__) -# Try to import tos module, and provide helpful error message if it fails -try: - import tos -except ImportError as e: - logger.error( - "Failed to import 'tos' module. Please install it using: pip install tos\n" - ) - raise ImportError( - "Missing 'tos' module. Please install it using: pip install tos\n" - ) from e - - -class TOSConfig(BaseModel): - region: str = Field( - default_factory=lambda: getenv("DATABASE_TOS_REGION"), - description="TOS region", - ) - ak: str = Field( - default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"), - description="Volcengine access key", - ) - sk: str = Field( - default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"), - description="Volcengine secret key", - ) - bucket_name: str = Field( - default_factory=lambda: getenv("DATABASE_TOS_BUCKET"), - description="TOS bucket name", - ) - - -class VeTOS(BaseModel): - config: TOSConfig = Field(default_factory=TOSConfig) - - def model_post_init(self, __context: Any) -> None: + +class VeTOS: + def __init__( + self, + ak: str = "", + sk: str = "", + region: str = "cn-beijing", + bucket_name: str = DEFAULT_TOS_BUCKET_NAME, + ) -> None: + self.ak = ak if ak else os.getenv("VOLCENGINE_ACCESS_KEY", "") + self.sk = sk if sk else os.getenv("VOLCENGINE_SECRET_KEY", "") + self.region = region + self.bucket_name = bucket_name + + try: + import tos + except ImportError as e: + logger.error( + "Failed to import 'tos' module. Please install it using: pip install tos\n" + ) + raise ImportError( + "Missing 'tos' module. Please install it using: pip install tos\n" + ) from e + + self._client = None try: self._client = tos.TosClientV2( - self.config.ak, - self.config.sk, - endpoint=f"tos-{self.config.region}.volces.com", - region=self.config.region, + ak=self.ak, + sk=self.sk, + endpoint=f"tos-{self.region}.volces.com", + region=self.region, ) - logger.info("Connected to TOS successfully.") + logger.info("Init TOS client.") except Exception as e: logger.error(f"Client initialization failed:{e}") - self._client = None def _refresh_client(self): try: if self._client: self._client.close() self._client = tos.TosClientV2( - self.config.ak, - self.config.sk, - endpoint=f"tos-{self.config.region}.volces.com", - region=self.config.region, + self.ak, + self.sk, + endpoint=f"tos-{self.region}.volces.com", + region=self.region, ) logger.info("refreshed client successfully.") except Exception as e: @@ -93,19 +85,17 @@ def create_bucket(self) -> bool: logger.error("TOS client is not initialized") return False try: - self._client.head_bucket(self.config.bucket_name) - logger.info(f"Bucket {self.config.bucket_name} already exists") + self._client.head_bucket(self.bucket_name) + logger.info(f"Bucket {self.bucket_name} already exists") except tos.exceptions.TosServerError as e: if e.status_code == 404: try: self._client.create_bucket( - bucket=self.config.bucket_name, + bucket=self.bucket_name, storage_class=tos.StorageClassType.Storage_Class_Standard, - acl=tos.ACLType.ACL_Public_Read, # 公开读 - ) - logger.info( - f"Bucket {self.config.bucket_name} created successfully" + acl=tos.ACLType.ACL_Public_Read, ) + logger.info(f"Bucket {self.bucket_name} created successfully") self._refresh_client() except Exception as create_error: logger.error(f"Bucket creation failed: {str(create_error)}") @@ -117,7 +107,7 @@ def create_bucket(self) -> bool: logger.error(f"Bucket check failed: {str(e)}") return False - # 确保在所有路径上返回布尔值 + # ensure return bool type return self._set_cors_rules() def _set_cors_rules(self) -> bool: @@ -131,14 +121,12 @@ def _set_cors_rules(self) -> bool: allowed_headers=["*"], max_age_seconds=1000, ) - self._client.put_bucket_cors(self.config.bucket_name, [rule]) - logger.info( - f"CORS rules for bucket {self.config.bucket_name} set successfully" - ) + self._client.put_bucket_cors(self.bucket_name, [rule]) + logger.info(f"CORS rules for bucket {self.bucket_name} set successfully") return True except Exception as e: logger.error( - f"Failed to set CORS rules for bucket {self.config.bucket_name}: {str(e)}" + f"Failed to set CORS rules for bucket {self.bucket_name}: {str(e)}" ) return False @@ -155,7 +143,9 @@ def build_tos_url( timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3] object_key: str = f"{app_name}-{user_id}-{session_id}/{timestamp}-{file_name}" - tos_url: str = f"https://{self.config.bucket_name}.tos-{self.config.region}.volces.com/{object_key}" + tos_url: str = ( + f"https://{self.bucket_name}.tos-{self.region}.volces.com/{object_key}" + ) return object_key, tos_url @@ -182,7 +172,7 @@ def _do_upload_bytes(self, object_key: str, data: bytes) -> None: if not self.create_bucket(): return self._client.put_object( - bucket=self.config.bucket_name, key=object_key, content=data + bucket=self.bucket_name, key=object_key, content=data ) logger.debug(f"Upload success, object_key: {object_key}") self._close() @@ -199,7 +189,7 @@ def _do_upload_file(self, object_key: str, file_path: str) -> None: if not self.create_bucket(): return self._client.put_object_from_file( - bucket=self.config.bucket_name, key=object_key, file_path=file_path + bucket=self.bucket_name, key=object_key, file_path=file_path ) self._close() logger.debug(f"Upload success, object_key: {object_key}") @@ -215,7 +205,7 @@ def download(self, object_key: str, save_path: str) -> bool: logger.error("TOS client is not initialized") return False try: - object_stream = self._client.get_object(self.config.bucket_name, object_key) + object_stream = self._client.get_object(self.bucket_name, object_key) save_dir = os.path.dirname(save_path) if save_dir and not os.path.exists(save_dir): diff --git a/veadk/runner.py b/veadk/runner.py index ddd5f8e0..a2b9d33b 100644 --- a/veadk/runner.py +++ b/veadk/runner.py @@ -27,12 +27,11 @@ from veadk.agents.loop_agent import LoopAgent from veadk.agents.parallel_agent import ParallelAgent from veadk.agents.sequential_agent import SequentialAgent -from veadk.config import getenv from veadk.evaluation import EvalSetRecorder from veadk.memory.short_term_memory import ShortTermMemory from veadk.types import MediaMessage from veadk.utils.logger import get_logger -from veadk.utils.misc import read_png_to_bytes +from veadk.utils.misc import getenv, read_png_to_bytes logger = get_logger(__name__) diff --git a/veadk/tools/builtin_tools/vesearch.py b/veadk/tools/builtin_tools/vesearch.py index d707ff4b..9937a0c9 100644 --- a/veadk/tools/builtin_tools/vesearch.py +++ b/veadk/tools/builtin_tools/vesearch.py @@ -14,7 +14,7 @@ import requests -from veadk.config import getenv +from veadk.config import getenv, settings def vesearch(query: str) -> str: @@ -26,7 +26,7 @@ def vesearch(query: str) -> str: Returns: Summarized search results. """ - api_key = getenv("TOOL_VESEARCH_API_KEY") + api_key = settings.tool.vesearch.api_key bot_id = str(getenv("TOOL_VESEARCH_ENDPOINT")) if api_key == "": diff --git a/veadk/tools/builtin_tools/video_generate.py b/veadk/tools/builtin_tools/video_generate.py index b4ff88df..d89097c3 100644 --- a/veadk/tools/builtin_tools/video_generate.py +++ b/veadk/tools/builtin_tools/video_generate.py @@ -12,18 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict -from google.adk.tools import ToolContext -from volcenginesdkarkruntime import Ark -from veadk.config import getenv +import json import time import traceback -import json -from veadk.version import VERSION +from typing import Dict, cast + +from google.adk.tools import ToolContext from opentelemetry import trace from opentelemetry.trace import Span +from volcenginesdkarkruntime import Ark +from volcenginesdkarkruntime.types.content_generation.create_task_content_param import ( + CreateTaskContentParam, +) +from veadk.config import getenv from veadk.utils.logger import get_logger +from veadk.version import VERSION logger = get_logger(__name__) @@ -47,13 +51,16 @@ async def generate(prompt, first_frame_image=None, last_frame_image=None): logger.debug("first frame generation") response = client.content_generation.tasks.create( model=getenv("MODEL_VIDEO_NAME"), - content=[ - {"type": "text", "text": prompt}, - { - "type": "image_url", - "image_url": {"url": first_frame_image}, - }, - ], + content=cast( + list[CreateTaskContentParam], # avoid IDE warning + [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": {"url": first_frame_image}, + }, + ], + ), ) else: logger.debug("last frame generation") @@ -277,13 +284,13 @@ async def video_generate(params: list, tool_context: ToolContext) -> Dict: def add_span_attributes( span: Span, tool_context: ToolContext, - input_part: dict = None, - output_part: dict = None, - input_tokens: int = None, - output_tokens: int = None, - total_tokens: int = None, - request_model: str = None, - response_model: str = None, + input_part: dict | None = None, + output_part: dict | None = None, + input_tokens: int | None = None, + output_tokens: int | None = None, + total_tokens: int | None = None, + request_model: str | None = None, + response_model: str | None = None, ): try: # common attributes diff --git a/veadk/tracing/telemetry/exporters/apmplus_exporter.py b/veadk/tracing/telemetry/exporters/apmplus_exporter.py index 32c53d94..0f94c350 100644 --- a/veadk/tracing/telemetry/exporters/apmplus_exporter.py +++ b/veadk/tracing/telemetry/exporters/apmplus_exporter.py @@ -28,7 +28,7 @@ from pydantic import BaseModel, Field from typing_extensions import override -from veadk.config import getenv +from veadk.config import settings from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter from veadk.utils.logger import get_logger @@ -102,19 +102,13 @@ def record(self, llm_request: LlmRequest, llm_response: LlmResponse) -> None: class APMPlusExporterConfig(BaseModel): endpoint: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_APMPLUS_ENDPOINT", - "http://apmplus-cn-beijing.volces.com:4317", - ), + default_factory=lambda: settings.apmplus_config.otel_exporter_endpoint, ) app_key: str = Field( - default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY"), + default_factory=lambda: settings.apmplus_config.otel_exporter_api_key, ) service_name: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_APMPLUS_SERVICE_NAME", - "veadk_tracing_service", - ), + default_factory=lambda: settings.apmplus_config.otel_exporter_service_name, description="Service name shown in APMPlus frontend.", ) @@ -123,6 +117,7 @@ class APMPlusExporter(BaseExporter): config: APMPlusExporterConfig = Field(default_factory=APMPlusExporterConfig) def model_post_init(self, context: Any) -> None: + print(self.config) headers = { "x-byteapm-appkey": self.config.app_key, } diff --git a/veadk/tracing/telemetry/exporters/cozeloop_exporter.py b/veadk/tracing/telemetry/exporters/cozeloop_exporter.py index ef4e4ed8..2d5b9b7e 100644 --- a/veadk/tracing/telemetry/exporters/cozeloop_exporter.py +++ b/veadk/tracing/telemetry/exporters/cozeloop_exporter.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field from typing_extensions import override -from veadk.config import getenv +from veadk.config import settings from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter from veadk.utils.logger import get_logger @@ -28,18 +28,13 @@ class CozeloopExporterConfig(BaseModel): endpoint: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT", - "https://api.coze.cn/v1/loop/opentelemetry/v1/traces", - ), + default_factory=lambda: settings.cozeloop_config.otel_exporter_endpoint, ) space_id: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME" - ), + default_factory=lambda: settings.cozeloop_config.otel_exporter_space_id, ) token: str = Field( - default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY"), + default_factory=lambda: settings.cozeloop_config.otel_exporter_api_key, ) diff --git a/veadk/tracing/telemetry/exporters/tls_exporter.py b/veadk/tracing/telemetry/exporters/tls_exporter.py index 29829205..8d2e4cea 100644 --- a/veadk/tracing/telemetry/exporters/tls_exporter.py +++ b/veadk/tracing/telemetry/exporters/tls_exporter.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field from typing_extensions import override -from veadk.config import getenv +from veadk.config import getenv, settings from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter from veadk.utils.logger import get_logger @@ -28,19 +28,13 @@ class TLSExporterConfig(BaseModel): endpoint: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT", - "https://tls-cn-beijing.volces.com:4318/v1/traces", - ), + default_factory=lambda: settings.tls_config.otel_exporter_endpoint, ) region: str = Field( - default_factory=lambda: getenv( - "OBSERVABILITY_OPENTELEMETRY_TLS_REGION", - "cn-beijing", - ), + default_factory=lambda: settings.tls_config.otel_exporter_region, ) topic_id: str = Field( - default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME"), + default_factory=lambda: settings.tls_config.otel_exporter_topic_id, ) access_key: str = Field(default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")) secret_key: str = Field(default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")) diff --git a/veadk/utils/logger.py b/veadk/utils/logger.py index c3a749c6..d8350d5b 100644 --- a/veadk/utils/logger.py +++ b/veadk/utils/logger.py @@ -16,7 +16,7 @@ from loguru import logger -from veadk.config import getenv +from veadk.utils.misc import getenv def filter_log(): diff --git a/veadk/utils/misc.py b/veadk/utils/misc.py index f6e35648..9761950e 100644 --- a/veadk/utils/misc.py +++ b/veadk/utils/misc.py @@ -14,12 +14,14 @@ import importlib.util import json +import os import sys import time import types from typing import Any, Dict, List, MutableMapping, Tuple import requests +from yaml import safe_load def read_file(file_path): @@ -100,3 +102,49 @@ def safe_json_serialize(obj) -> str: ) except (TypeError, OverflowError): return "" + + +def getenv( + env_name: str, default_value: Any = "", allow_false_values: bool = False +) -> str: + """ + Get environment variable. + + Args: + env_name (str): The name of the environment variable. + default_value (str): The default value of the environment variable. + allow_false_values (bool, optional): Whether to allow the environment variable to be None or false values. Defaults to False. + + Returns: + str: The value of the environment variable. + """ + value = os.getenv(env_name, default_value) + + if allow_false_values: + return value + + if value: + return value + else: + raise ValueError( + f"The environment variable `{env_name}` not exists. Please set this in your environment variable or config.yaml." + ) + + +def set_envs(config_yaml_path: str) -> tuple[dict, dict]: + with open(config_yaml_path, "r", encoding="utf-8") as yaml_file: + config_dict = safe_load(yaml_file) + + flatten_config_dict = flatten_dict(config_dict) + + for k, v in flatten_config_dict.items(): + veadk_environments = {} + + k = k.upper() + if k in os.environ: + veadk_environments[k] = os.environ[k] + continue + veadk_environments[k] = str(v) + os.environ[k] = str(v) + + return config_dict, veadk_environments diff --git a/veadk/utils/volcengine_sign.py b/veadk/utils/volcengine_sign.py index 34d7089d..4f945d76 100644 --- a/veadk/utils/volcengine_sign.py +++ b/veadk/utils/volcengine_sign.py @@ -169,6 +169,7 @@ def ve_request( region: str, host: str, content_type: str = "application/json", + header: dict = {}, ): # response_body = request("Get", datetime.datetime.utcnow(), {}, {}, AK, SK, "ListUsers", None) # print(response_body) @@ -195,7 +196,7 @@ def ve_request( try: response_body = request( - "POST", now, {}, {}, AK, SK, action, json.dumps(request_body) + "POST", now, {}, header, AK, SK, action, json.dumps(request_body) ) return response_body except Exception as e: From d2b51a7e1ab2b28b17428bdba9a1dfc3109ad458 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 10 Sep 2025 10:48:54 +0800 Subject: [PATCH 3/6] fix config path and envs --- veadk/config.py | 8 ++++++-- veadk/configs/database_configs.py | 11 ++--------- veadk/configs/tracing_configs.py | 12 +++++++++++- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/veadk/config.py b/veadk/config.py index 6ceeb076..e3344a03 100644 --- a/veadk/config.py +++ b/veadk/config.py @@ -21,14 +21,18 @@ from veadk.configs.database_configs import ( MysqlConfig, OpensearchConfig, - PrometheusConfig, RedisConfig, TOSConfig, VikingKnowledgebaseConfig, ) from veadk.configs.model_configs import ModelConfig from veadk.configs.tool_configs import BuiltinToolConfigs, PromptPilotConfig -from veadk.configs.tracing_configs import APMPlusConfig, CozeloopConfig, TLSConfig +from veadk.configs.tracing_configs import ( + APMPlusConfig, + CozeloopConfig, + PrometheusConfig, + TLSConfig, +) from veadk.utils.misc import set_envs diff --git a/veadk/configs/database_configs.py b/veadk/configs/database_configs.py index d8879cb5..86d8300f 100644 --- a/veadk/configs/database_configs.py +++ b/veadk/configs/database_configs.py @@ -15,7 +15,6 @@ import os from functools import cached_property -from pydantic import BaseModel from pydantic_settings import BaseSettings, SettingsConfigDict from veadk.consts import DEFAULT_TOS_BUCKET_NAME @@ -60,15 +59,9 @@ class RedisConfig(BaseSettings): db: int = 0 -class PrometheusConfig(BaseSettings): - pushgateway_url: str = "" +class VikingKnowledgebaseConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_VIKING_") - pushgateway_username: str = "" - - pushgateway_password: str = "" - - -class VikingKnowledgebaseConfig(BaseModel): project: str = "default" """User project in Volcengine console web.""" diff --git a/veadk/configs/tracing_configs.py b/veadk/configs/tracing_configs.py index a53c5d72..2b913d73 100644 --- a/veadk/configs/tracing_configs.py +++ b/veadk/configs/tracing_configs.py @@ -16,7 +16,7 @@ from functools import cached_property from pydantic import Field -from pydantic_settings import BaseSettings +from pydantic_settings import BaseSettings, SettingsConfigDict from veadk.auth.veauth.apmplus_veauth import APMPlusVeAuth from veadk.consts import ( @@ -81,3 +81,13 @@ def otel_exporter_topic_id(self) -> str: or VeTLS().get_trace_topic_id() ) return _topic_id + + +class PrometheusConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="OBSERVABILITY_PROMETHEUS_") + + pushgateway_url: str = "" + + pushgateway_username: str = "" + + pushgateway_password: str = "" From e7942719dfd6f3ecf61d9d31a48529e88a8f01c9 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 10 Sep 2025 18:22:27 +0800 Subject: [PATCH 4/6] fix tls bugs --- .../{{cookiecutter.local_dir_name}}/deploy.py | 4 +- veadk/integrations/ve_tls/utils.py | 6 +- veadk/integrations/ve_tls/ve_tls.py | 56 ++++++++----------- .../telemetry/exporters/apmplus_exporter.py | 1 - 4 files changed, 28 insertions(+), 39 deletions(-) diff --git a/veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/deploy.py b/veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/deploy.py index b3a77492..7ceb5bf4 100644 --- a/veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/deploy.py +++ b/veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/deploy.py @@ -95,8 +95,8 @@ async def main(): message = "How is the weather like in Beijing?" print(f"Test message: {message}") - await _send_msg_with_a2a(cloud_app=cloud_app, message=message) - await _send_msg_with_mcp(cloud_app=cloud_app, message=message) + # await _send_msg_with_a2a(cloud_app=cloud_app, message=message) + # await _send_msg_with_mcp(cloud_app=cloud_app, message=message) if __name__ == "__main__": diff --git a/veadk/integrations/ve_tls/utils.py b/veadk/integrations/ve_tls/utils.py index 33009400..a18a2784 100644 --- a/veadk/integrations/ve_tls/utils.py +++ b/veadk/integrations/ve_tls/utils.py @@ -80,7 +80,7 @@ def __prepare_request( return request -async def ve_tls_request( +def ve_tls_request( client: TLSService, api: str, params: dict | None = None, @@ -100,10 +100,10 @@ async def ve_tls_request( url = request.build() - async with httpx.AsyncClient() as session: + with httpx.Client() as session: while True: try: - response = await session.request( + response = session.request( method=method, url=url, headers=request.headers, diff --git a/veadk/integrations/ve_tls/ve_tls.py b/veadk/integrations/ve_tls/ve_tls.py index 3828eda6..93f06dc4 100644 --- a/veadk/integrations/ve_tls/ve_tls.py +++ b/veadk/integrations/ve_tls/ve_tls.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import asyncio import os +from volcengine.tls.TLSService import TLSService + from veadk.consts import DEFAULT_TLS_LOG_PROJECT_NAME, DEFAULT_TLS_TRACING_INSTANCE_NAME from veadk.integrations.ve_tls.utils import ve_tls_request from veadk.utils.logger import get_logger -from volcengine.tls.TLSService import TLSService logger = get_logger(__name__) @@ -63,13 +63,11 @@ def get_project_id_by_name(self, project_name: str) -> str: try: res = None - res = asyncio.run( - ve_tls_request( - client=self._client, - api="DescribeProjects", - body=request_body, - method="GET", - ) + res = ve_tls_request( + client=self._client, + api="DescribeProjects", + body=request_body, + method="GET", ) projects = res["Projects"] for project in projects: @@ -97,10 +95,8 @@ def create_log_project(self, project_name: str) -> str: "Tags": [{"Key": "provider", "Value": "VeADK"}], } try: - res = asyncio.run( - ve_tls_request( - client=self._client, api="CreateProject", body=request_body - ) + res = ve_tls_request( + client=self._client, api="CreateProject", body=request_body ) if res["ErrorCode"] == "ProjectAlreadyExist": @@ -122,13 +118,11 @@ def get_trace_instance_by_name(self, log_project_id: str, trace_instance_name: s "TraceInstanceName": trace_instance_name, } try: - res = asyncio.run( - ve_tls_request( - client=self._client, - api="DescribeTraceInstances", - body=request_body, - method="GET", - ) + res = ve_tls_request( + client=self._client, + api="DescribeTraceInstances", + body=request_body, + method="GET", ) for instance in res["TraceInstances"]: @@ -156,12 +150,10 @@ def create_tracing_instance(self, log_project_id: str, trace_instance_name: str) try: res = None - res = asyncio.run( - ve_tls_request( - client=self._client, - api="CreateTraceInstance", - body=request_body, - ) + res = ve_tls_request( + client=self._client, + api="CreateTraceInstance", + body=request_body, ) if res["ErrorCode"] == "TopicAlreadyExist": @@ -173,13 +165,11 @@ def create_tracing_instance(self, log_project_id: str, trace_instance_name: str) ) # after creation, get the trace instance details - res = asyncio.run( - ve_tls_request( - client=self._client, - api="DescribeTraceInstance", - body={"TraceInstanceID": res["TraceInstanceID"]}, - method="GET", - ) + res = ve_tls_request( + client=self._client, + api="DescribeTraceInstance", + body={"TraceInstanceID": res["TraceInstanceID"]}, + method="GET", ) return res diff --git a/veadk/tracing/telemetry/exporters/apmplus_exporter.py b/veadk/tracing/telemetry/exporters/apmplus_exporter.py index 0f94c350..72bf4efe 100644 --- a/veadk/tracing/telemetry/exporters/apmplus_exporter.py +++ b/veadk/tracing/telemetry/exporters/apmplus_exporter.py @@ -117,7 +117,6 @@ class APMPlusExporter(BaseExporter): config: APMPlusExporterConfig = Field(default_factory=APMPlusExporterConfig) def model_post_init(self, context: Any) -> None: - print(self.config) headers = { "x-byteapm-appkey": self.config.app_key, } From adaac82dc520f74f7a45d9bc6182827a31662093 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 10 Sep 2025 18:32:27 +0800 Subject: [PATCH 5/6] add some logs --- veadk/auth/veauth/vesearch_veauth.py | 2 +- veadk/cli/cli_prompt.py | 4 ++-- veadk/integrations/ve_tls/ve_tls.py | 2 ++ 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/veadk/auth/veauth/vesearch_veauth.py b/veadk/auth/veauth/vesearch_veauth.py index 920bd001..4ca2010a 100644 --- a/veadk/auth/veauth/vesearch_veauth.py +++ b/veadk/auth/veauth/vesearch_veauth.py @@ -35,7 +35,7 @@ def __init__( @override def _fetch_token(self): - logger.info("Fetching Vesearch token from Volcengine Top Gateway.") + logger.info("Fetching VeSearch token ...") res = ve_request( request_body={"biz_scene": "search_agent", "page": 1, "rows": 10}, diff --git a/veadk/cli/cli_prompt.py b/veadk/cli/cli_prompt.py index 5f05b6cb..f621da28 100644 --- a/veadk/cli/cli_prompt.py +++ b/veadk/cli/cli_prompt.py @@ -19,8 +19,8 @@ @click.option( "--path", default=".", help="Agent file path with global variable `agent=...`" ) -@click.option("--feedback", default=None, help="Suggestions for prompt optimization") -@click.option("--api-key", default=None, help="API Key of PromptPilot") +@click.option("--feedback", default="", help="Suggestions for prompt optimization") +@click.option("--api-key", default="", help="API Key of PromptPilot") @click.option( "--model-name", default="doubao-1.5-pro-32k-250115", diff --git a/veadk/integrations/ve_tls/ve_tls.py b/veadk/integrations/ve_tls/ve_tls.py index 93f06dc4..f63dbf64 100644 --- a/veadk/integrations/ve_tls/ve_tls.py +++ b/veadk/integrations/ve_tls/ve_tls.py @@ -195,4 +195,6 @@ def get_trace_topic_id(self): if not instance: raise ValueError("None instance") + logger.info(f"Fetched trace topic id: {instance['TraceTopicId']}") + return instance["TraceTopicId"] From 6e1b4773d720e9e8b191fca95695dfa6eb21bf4e Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 10 Sep 2025 19:14:38 +0800 Subject: [PATCH 6/6] fix envs --- veadk/utils/misc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/veadk/utils/misc.py b/veadk/utils/misc.py index 9761950e..ca1f5df4 100644 --- a/veadk/utils/misc.py +++ b/veadk/utils/misc.py @@ -137,10 +137,10 @@ def set_envs(config_yaml_path: str) -> tuple[dict, dict]: flatten_config_dict = flatten_dict(config_dict) + veadk_environments = {} for k, v in flatten_config_dict.items(): - veadk_environments = {} - k = k.upper() + if k in os.environ: veadk_environments[k] = os.environ[k] continue