veadk-python 0.2.8__tar.gz → 0.2.10__tar.gz

{veadk_python-0.2.8 → veadk_python-0.2.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.2.8
+Version: 0.2.10
 Summary: Volcengine agent development kit, integrations with Volcengine cloud services.
 Author-email: Yaozheng Fang <fangyozheng@gmail.com>, Guodong Li <cu.eric.lee@gmail.com>, Zhi Han <sliverydayday@gmail.com>, Meng Wang <mengwangwm@gmail.com>
 License:                                  Apache License
@@ -218,9 +218,9 @@ Requires-Dist: opentelemetry-exporter-otlp>=1.35.0
 Requires-Dist: opentelemetry-instrumentation-logging>=0.56b0
 Requires-Dist: wrapt>=1.17.2
 Requires-Dist: openai<1.100
-Requires-Dist: volcengine-python-sdk>=4.0.3
+Requires-Dist: volcengine-python-sdk>=4.0.18
 Requires-Dist: volcengine>=1.0.193
-Requires-Dist: agent-pilot-sdk>=0.0.9
+Requires-Dist: agent-pilot-sdk>=0.1.2
 Requires-Dist: fastmcp>=2.11.3
 Requires-Dist: cookiecutter>=2.6.0
 Requires-Dist: omegaconf>=2.3.0
@@ -235,12 +235,12 @@ Provides-Extra: extensions
 Requires-Dist: redis>=5.0; extra == "extensions"
 Requires-Dist: tos>=2.8.4; extra == "extensions"
 Requires-Dist: llama-index-vector-stores-redis>=0.6.1; extra == "extensions"
-Requires-Dist: mcp-server-vikingdb-memory; extra == "extensions"
 Provides-Extra: database
 Requires-Dist: redis>=5.0; extra == "database"
 Requires-Dist: pymysql>=1.1.1; extra == "database"
 Requires-Dist: volcengine>=1.0.193; extra == "database"
 Requires-Dist: tos>=2.8.4; extra == "database"
+Requires-Dist: mem0ai==0.1.118; extra == "database"
 Provides-Extra: eval
 Requires-Dist: prometheus-client>=0.22.1; extra == "eval"
 Requires-Dist: deepeval>=3.2.6; extra == "eval"
@@ -270,6 +270,17 @@ A [tutorial](https://github.com/volcengine/veadk-python/blob/main/veadk_tutorial
 
 ## Installation
 
+### From PyPI
+
+```python
+pip install veadk-python
+
+# install extensions
+pip install veadk-python[extensions]
+```
+
+### Build from source
+
 We use `uv` to build this project ([how-to-install-uv](https://docs.astral.sh/uv/getting-started/installation/)).
 
 ```bash

{veadk_python-0.2.8 → veadk_python-0.2.10}/README.md

@@ -15,6 +15,17 @@ A [tutorial](https://github.com/volcengine/veadk-python/blob/main/veadk_tutorial
 
 ## Installation
 
+### From PyPI
+
+```python
+pip install veadk-python
+
+# install extensions
+pip install veadk-python[extensions]
+```
+
+### Build from source
+
 We use `uv` to build this project ([how-to-install-uv](https://docs.astral.sh/uv/getting-started/installation/)).
 
 ```bash

{veadk_python-0.2.8 → veadk_python-0.2.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.8"
+version = "0.2.10"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,9 +22,9 @@ dependencies = [
     "opentelemetry-instrumentation-logging>=0.56b0",
     "wrapt>=1.17.2", # For patching built-in functions
     "openai<1.100", # For fix https://github.com/BerriAI/litellm/issues/13710
-    "volcengine-python-sdk>=4.0.3", # For Volcengine API
-    "volcengine>=1.0.193",          # For Volcengine sign
-    "agent-pilot-sdk>=0.0.9", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
+    "volcengine-python-sdk>=4.0.18", # For Volcengine API
+    "volcengine>=1.0.193", # For Volcengine sign
+    "agent-pilot-sdk>=0.1.2", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
     "fastmcp>=2.11.3", # For running MCP
     "cookiecutter>=2.6.0", # For cloud deploy
     "omegaconf>=2.3.0", # For agent builder
@@ -45,13 +45,13 @@ extensions = [
     "redis>=5.0",                 # For Redis database
     "tos>=2.8.4",                   # For TOS storage and Viking DB
     "llama-index-vector-stores-redis>=0.6.1",
-    "mcp-server-vikingdb-memory",
 ]
 database = [
     "redis>=5.0",                 # For Redis database
     "pymysql>=1.1.1",               # For MySQL database
     "volcengine>=1.0.193",          # For Viking DB
     "tos>=2.8.4",                   # For TOS storage and Viking DB
+    "mem0ai==0.1.118",              # For mem0
 ]
 eval = [
     "prometheus-client>=0.22.1",    # For exporting data to Prometheus pushgateway
@@ -91,6 +91,3 @@ exclude = [
     "veadk/integrations/ve_faas/template/*",
     "veadk/integrations/ve_faas/web_template/*"
 ]
-
-[tool.uv.sources]
-mcp-server-vikingdb-memory = { git = "https://github.com/volcengine/mcp-server", subdirectory = "server/mcp_server_vikingdb_memory" }

veadk_python-0.2.10/veadk/a2a/remote_ve_agent.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 json
+from typing import Literal
+
+import httpx
+import requests
+from a2a.types import AgentCard
+from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
+
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json"
+
+
+def _convert_agent_card_dict_to_obj(agent_card_dict: dict) -> AgentCard:
+    agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+    agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+    return agent_card_object
+
+
+class RemoteVeAgent(RemoteA2aAgent):
+    """Connect to remote agent on VeFaaS platform."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        auth_token: str | None = None,
+        auth_method: Literal["header", "querystring"] | None = None,
+    ):
+        if not auth_token:
+            agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
+            # replace agent_card_url with actual host
+            agent_card_dict["url"] = url
+
+            agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+
+            logger.debug(f"Agent card of {name}: {agent_card_object}")
+            super().__init__(name=name, agent_card=agent_card_object)
+        else:
+            if auth_method == "header":
+                headers = {"Authorization": f"Bearer {auth_token}"}
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH, headers=headers
+                ).json()
+                agent_card_dict["url"] = url
+
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, headers=headers, timeout=600
+                )
+
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            elif auth_method == "querystring":
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH + f"?token={auth_token}"
+                ).json()
+                agent_card_dict["url"] = url
+
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, params={"token": auth_token}, timeout=600
+                )
+
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported auth method {auth_method}, use `header` or `querystring` instead."
+                )

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/agent.py

@@ -84,6 +84,9 @@ class Agent(LlmAgent):
     knowledgebase: Optional[KnowledgeBase] = None
     """The knowledgebase provided to agent."""
 
+    short_term_memory: Optional[ShortTermMemory] = None
+    """The short term memory provided to agent."""
+
     long_term_memory: Optional[LongTermMemory] = None
     """The long term memory provided to agent.
 
@@ -138,6 +141,9 @@ class Agent(LlmAgent):
         if self.long_term_memory is not None:
             from google.adk.tools import load_memory
 
+            if not load_memory.custom_metadata:
+                load_memory.custom_metadata = {}
+            load_memory.custom_metadata["backend"] = self.long_term_memory.backend
             self.tools.append(load_memory)
 
         logger.info(f"VeADK version: {VERSION}")

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/agent_builder.py

@@ -52,9 +52,8 @@ class AgentBuilder:
         tools = []
         if agent_config.get("tools", []):
             for tool in agent_config["tools"]:
-                module_name = tool["module"]
-                func_name = tool["func"]
-
+                name = tool["name"]
+                module_name, func_name = name.rsplit(".", 1)
                 module = importlib.import_module(module_name)
                 func = getattr(module, func_name)
 

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/cli/cli.py

@@ -18,6 +18,7 @@ import click
 from veadk.cli.cli_deploy import deploy
 from veadk.cli.cli_eval import eval
 from veadk.cli.cli_init import init
+from veadk.cli.cli_kb import kb
 from veadk.cli.cli_pipeline import pipeline
 from veadk.cli.cli_prompt import prompt
 from veadk.cli.cli_web import web
@@ -39,6 +40,7 @@ veadk.add_command(prompt)
 veadk.add_command(web)
 veadk.add_command(pipeline)
 veadk.add_command(eval)
+veadk.add_command(kb)
 
 if __name__ == "__main__":
     veadk()

veadk_python-0.2.10/veadk/cli/cli_kb.py

@@ -0,0 +1,75 @@
+# 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 pathlib import Path
+from typing import Literal
+
+import click
+
+
+@click.command()
+@click.option(
+    "--backend",
+    type=click.Choice(
+        ["local", "opensearch", "viking", "redis"],
+        case_sensitive=False,
+    ),
+    required=True,
+)
+@click.option(
+    "--app_name",
+    default="",
+    help="`app_name` for init your knowledgebase",
+)
+@click.option(
+    "--index",
+    default="",
+    help="Knowledgebase index",
+)
+@click.option(
+    "--path",
+    required=True,
+    help="Knowledge file or directory path",
+)
+def add(
+    backend: Literal["local", "opensearch", "viking", "redis"],
+    app_name: str,
+    index: str,
+    path: str,
+):
+    """Add files to knowledgebase"""
+    _path = Path(path)
+    assert _path.exists(), f"Path {path} not exists. Please check your input."
+
+    from veadk.knowledgebase import KnowledgeBase
+
+    knowledgebase = KnowledgeBase(backend=backend, app_name=app_name, index=index)
+
+    if _path.is_file():
+        knowledgebase.add_from_files(files=[path])
+    elif _path.is_dir():
+        knowledgebase.add_from_directory(directory=path)
+    else:
+        raise RuntimeError(
+            "Unsupported knowledgebase file type, only support a single file and a directory."
+        )
+
+
+@click.group()
+def kb():
+    """VeADK Knowledgebase management"""
+    pass
+
+
+kb.add_command(add)

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/cli/cli_prompt.py

@@ -21,12 +21,15 @@ import click
 )
 @click.option("--feedback", default="", help="Suggestions for prompt optimization")
 @click.option("--api-key", default="", help="API Key of PromptPilot")
+@click.option("--workspace-id", default="", help="Workspace ID of PromptPilot")
 @click.option(
     "--model-name",
     default="doubao-1.5-pro-32k-250115",
     help="Model name for prompt optimization",
 )
-def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None:
+def prompt(
+    path: str, feedback: str, api_key: str, workspace_id: str, model_name: str
+) -> None:
     """Optimize agent system prompt from a local file."""
     from pathlib import Path
 
@@ -56,7 +59,11 @@ def prompt(path: str, feedback: str, api_key: str, model_name: str) -> None:
 
         if not api_key:
             api_key = settings.prompt_pilot.api_key
-        ve_prompt_pilot = VePromptPilot(api_key)
+
+        if not workspace_id:
+            raise ValueError("Please provide workspace_id for PromptPilot service.")
+
+        ve_prompt_pilot = VePromptPilot(api_key=api_key, workspace_id=workspace_id)
         ve_prompt_pilot.optimize(
             agents=agents, feedback=feedback, model_name=model_name
         )

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/cli/cli_web.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 
 import click
 
@@ -152,6 +153,8 @@ def web(host: str) -> None:
         eval_sets_manager: Any,
         eval_set_results_manager: Any,
         agents_dir: str,
+        extra_plugins: Optional[list[str]] = None,
+        **kwargs: Any,
     ):
         self.agent_loader = agent_loader
         self.artifact_service = artifact_service
@@ -162,6 +165,10 @@ def web(host: str) -> None:
         self.runners_to_clean = set()
         self.current_app_name_ref = SharedValue(value="")
         self.runner_dict = {}
+        self.extra_plugins = extra_plugins or []
+
+        for key, value in kwargs.items():
+            setattr(self, key, value)
 
         # parse VeADK memories
         short_term_memory, long_term_memory = _get_memory(module_path=agents_dir)

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/configs/database_configs.py

@@ -83,6 +83,15 @@ class RedisConfig(BaseSettings):
     """STS token for Redis auth, not supported yet."""
 
 
+class Mem0Config(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="DATABASE_MEM0_")
+
+    api_key: str = ""
+    """Mem0 API key"""
+
+    base_url: str = ""  # "https://api.mem0.ai/v1"
+
+
 class VikingKnowledgebaseConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_VIKING_")
 

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/consts.py

@@ -27,6 +27,8 @@ DEFAULT_MODEL_EXTRA_CONFIG = {
         "x-is-encrypted": getenv("MODEL_AGENT_ENCRYPTED", "true"),
         "veadk-source": "veadk",
         "veadk-version": VERSION,
+        "User-Agent": f"VeADK/{VERSION}",
+        "X-Client-Request-Id": f"veadk/{VERSION}",
     },
     "extra_body": {
         "caching": {
@@ -59,3 +61,8 @@ DEFAULT_TLS_TRACING_INSTANCE_NAME = "veadk-tracing"
 DEFAULT_TOS_BUCKET_NAME = "ark-tutorial"
 
 DEFAULT_COZELOOP_SPACE_NAME = "VeADK Space"
+
+DEFAULT_TEXT_TO_IMAGE_MODEL_NAME = "doubao-seedream-3-0-t2i-250415"
+DEFAULT_IMAGE_EDIT_MODEL_NAME = "doubao-seededit-3-0-i2i-250628"
+DEFAULT_VIDEO_MODEL_NAME = "doubao-seedance-1-0-pro-250528"
+DEFAULT_IMAGE_GENERATE_MODEL_NAME = "doubao-seedream-4-0-250828"

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/evaluation/adk_evaluator/adk_evaluator.py

@@ -24,6 +24,8 @@ from google.adk.evaluation.agent_evaluator import (
 )
 from google.adk.evaluation.eval_case import IntermediateData, Invocation
 from google.adk.evaluation.evaluator import EvalStatus
+from google.adk.evaluation.eval_set import EvalSet
+from typing import Optional
 from typing_extensions import override
 from veadk.evaluation.base_evaluator import BaseEvaluator
 from types import SimpleNamespace
@@ -52,7 +54,8 @@ class ADKEvaluator(BaseEvaluator):
     @override
     async def evaluate(
         self,
-        eval_set_file_path: str,
+        eval_set: Optional[EvalSet] = None,
+        eval_set_file_path: Optional[str] = None,
         eval_id: str = f"test_{formatted_timestamp()}",
         tool_score_threshold: float = 1.0,
         response_match_score_threshold: float = 0.8,
@@ -104,7 +107,7 @@ class ADKEvaluator(BaseEvaluator):
         # Iterate each test file and evaluate per-case, per-metric
         for test_file in test_files:
             # Build in-memory evaluation cases via BaseEvaluator from the provided file
-            self.build_eval_set(test_file)
+            self.build_eval_set(eval_set, test_file)
 
             evaluation_result_list = []
 

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/evaluation/base_evaluator.py

@@ -17,7 +17,7 @@ import json
 import time
 import uuid
 from abc import abstractmethod
-from typing import Any
+from typing import Any, Optional
 
 from google.adk import Runner
 from google.adk.evaluation.eval_set import EvalSet
@@ -210,33 +210,43 @@ class BaseEvaluator:
 
         return evalset
 
-    def build_eval_set(self, file_path: str):
+    def build_eval_set(
+        self, eval_set: Optional[EvalSet] = None, file_path: Optional[str] = None
+    ):
         """Generate evaluation data from a given file and assign it to the class attribute `invocation_list`."""
-        eval_case_data_list: list[EvalTestCase] = []
 
-        try:
-            with open(file_path, "r", encoding="utf-8") as f:
-                file_content = json.load(f)
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Invalid JSON format in file {file_path}: {e}")
-        except Exception as e:
-            raise ValueError(f"Error reading file {file_path}: {e}")
-
-        if isinstance(file_content, dict) and "eval_cases" in file_content:
-            eval_cases = self._build_eval_set_from_eval_json(file_path).eval_cases
-        elif (
-            isinstance(file_content, list)
-            and len(file_content) > 0
-            and all(
-                isinstance(span, dict) and "trace_id" in span for span in file_content
-            )
-        ):
-            eval_cases = self._build_eval_set_from_tracing_json(file_path).eval_cases
+        if eval_set is None and file_path is None:
+            raise ValueError("eval_set or file_path is required")
+        if eval_set:
+            eval_cases = eval_set.eval_cases
         else:
-            raise ValueError(
-                f"Unsupported file format in {file_path}. Please provide a valid file."
-            )
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    file_content = json.load(f)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON format in file {file_path}: {e}")
+            except Exception as e:
+                raise ValueError(f"Error reading file {file_path}: {e}")
+
+            if isinstance(file_content, dict) and "eval_cases" in file_content:
+                eval_cases = self._build_eval_set_from_eval_json(file_path).eval_cases
+            elif (
+                isinstance(file_content, list)
+                and len(file_content) > 0
+                and all(
+                    isinstance(span, dict) and "trace_id" in span
+                    for span in file_content
+                )
+            ):
+                eval_cases = self._build_eval_set_from_tracing_json(
+                    file_path
+                ).eval_cases
+            else:
+                raise ValueError(
+                    f"Unsupported file format in {file_path}. Please provide a valid file."
+                )
 
+        eval_case_data_list: list[EvalTestCase] = []
         for eval_case in eval_cases:
             eval_case_data = EvalTestCase(invocations=[])
             if eval_case.session_input:
@@ -384,8 +394,9 @@ class BaseEvaluator:
     @abstractmethod
     async def evaluate(
         self,
-        eval_set_file_path: str,
         metrics: list[Any],
+        eval_set: Optional[EvalSet],
+        eval_set_file_path: Optional[str],
         eval_id: str,
     ):
         """An abstract method for evaluation based on metrics。"""

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/evaluation/deepeval_evaluator/deepeval_evaluator.py

@@ -22,7 +22,8 @@ from deepeval.models import LocalModel
 from deepeval.test_case import LLMTestCase
 from deepeval.test_case.llm_test_case import ToolCall
 from typing_extensions import override
-
+from typing import Optional
+from google.adk.evaluation.eval_set import EvalSet
 from veadk.config import getenv
 from veadk.evaluation.base_evaluator import BaseEvaluator, EvalResultData, MetricResult
 from veadk.evaluation.types import EvalResultCaseData, EvalResultMetadata
@@ -77,13 +78,14 @@ class DeepevalEvaluator(BaseEvaluator):
     @override
     async def evaluate(
         self,
-        eval_set_file_path: str,
         metrics: list[BaseMetric],
+        eval_set: Optional[EvalSet] = None,
+        eval_set_file_path: Optional[str] = None,
         eval_id: str = f"test_{formatted_timestamp()}",
     ):
         """Target to Google ADK, we will use the same evaluation case format as Google ADK."""
         # Get evaluation data by parsing eval set file
-        self.build_eval_set(eval_set_file_path)
+        self.build_eval_set(eval_set, eval_set_file_path)
 
         # Get actual data by running agent
         logger.info("Start to run agent for actual data.")

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/integrations/ve_prompt_pilot/ve_prompt_pilot.py

@@ -16,7 +16,6 @@ from __future__ import annotations
 
 import agent_pilot as ap
 from agent_pilot.models import TaskType
-
 from veadk import Agent
 from veadk.prompts import prompt_optimization
 from veadk.utils.logger import get_logger
@@ -26,9 +25,15 @@ logger = get_logger(__name__)
 
 class VePromptPilot:
     def __init__(
-        self, api_key: str, path: str = "", task_id: str | None = None
+        self,
+        api_key: str,
+        workspace_id: str,
+        path: str = "",
+        task_id: str | None = None,
     ) -> None:
         self.api_key = api_key
+        self.workspace_id = workspace_id
+
         self.path = path
 
     def optimize(
@@ -57,12 +62,13 @@ class VePromptPilot:
             usage = None
             for chunk in ap.generate_prompt_stream(
                 task_description=task_description,
-                current_prompt=agent.instruction,
+                current_prompt=str(agent.instruction),
                 model_name=model_name,
                 task_type=TaskType.DIALOG,
                 temperature=1.0,
                 top_p=0.7,
                 api_key=self.api_key,
+                workspace_id=self.workspace_id,
             ):  # stream chunks of optimized prompt
                 # Process each chunk as it arrives
                 optimized_prompt += chunk.data.content if chunk.data else ""

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/integrations/ve_tls/utils.py

@@ -17,6 +17,7 @@ import json
 from typing import Any, Literal
 
 import httpx
+from veadk.utils.logger import get_logger
 from volcengine.ApiInfo import ApiInfo
 from volcengine.auth.SignerV4 import SignerV4
 from volcengine.tls.const import (
@@ -28,8 +29,6 @@ from volcengine.tls.const import (
 )
 from volcengine.tls.TLSService import TLSService
 
-from veadk.utils.logger import get_logger
-
 logger = get_logger(__name__)
 
 HEADER_API_VERSION = "x-tls-apiversion"

{veadk_python-0.2.8 → veadk_python-0.2.10}/veadk/integrations/ve_tls/ve_tls.py

@@ -28,9 +28,8 @@ class VeTLS:
         region: str = "cn-beijing",
     ):
         try:
-            from volcengine.tls.TLSService import TLSService
-
             from veadk.integrations.ve_tls.utils import ve_tls_request
+            from volcengine.tls.TLSService import TLSService
         except ImportError:
             raise ImportError(
                 "Please install volcengine SDK before init VeTLS: pip install volcengine"
@@ -162,11 +161,12 @@ class VeTLS:
                 client=self._client,
                 api="CreateTraceInstance",
                 body=request_body,
+                request_headers={"TraceTag": "veadk"},
             )
 
-            if res["ErrorCode"] == "TopicAlreadyExist":
+            if "ErrorCode" in res and res["ErrorCode"] == "TopicAlreadyExist":
                 logger.debug(
-                    f"Log project '{trace_instance_name}' already exists. Check its ID."
+                    f"Tracing instance '{trace_instance_name}' already exists. Check its ID."
                 )
                 return self.get_trace_instance_by_name(
                     log_project_id, trace_instance_name
@@ -176,10 +176,14 @@ class VeTLS:
             res = self._ve_tls_request(
                 client=self._client,
                 api="DescribeTraceInstance",
-                body={"TraceInstanceID": res["TraceInstanceID"]},
+                params={"TraceInstanceId": res["TraceInstanceId"]},
                 method="GET",
             )
 
+            logger.info(
+                f"Create tracing instance finished, tracing instance ID: {res['TraceInstanceId']}"
+            )
+
             return res
         except KeyError:
             raise ValueError(f"Failed to create tracing instance: {res}")