veadk-python 0.2.32__tar.gz → 0.5.5__tar.gz

{veadk_python-0.2.32 → veadk_python-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.2.32
+Version: 0.5.5
 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
@@ -208,40 +208,40 @@ License:                                  Apache License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pydantic-settings>=2.10.1
-Requires-Dist: a2a-sdk>=0.3.0
-Requires-Dist: deprecated>=1.2.18
-Requires-Dist: google-adk>=1.10.0
-Requires-Dist: google-adk<=1.19.0
+Requires-Dist: pydantic-settings==2.10.1
+Requires-Dist: a2a-sdk>=0.3.7
+Requires-Dist: deprecated==1.2.18
+Requires-Dist: google-adk>=1.19.0
 Requires-Dist: litellm>=1.74.3
-Requires-Dist: loguru>=0.7.3
-Requires-Dist: opentelemetry-exporter-otlp>=1.35.0
+Requires-Dist: loguru==0.7.3
+Requires-Dist: opentelemetry-exporter-otlp==1.37.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.33
+Requires-Dist: wrapt==1.17.2
+Requires-Dist: volcengine-python-sdk>=5.0.1
 Requires-Dist: volcengine>=1.0.193
-Requires-Dist: agent-pilot-sdk>=0.1.2
-Requires-Dist: fastmcp>=2.11.3
-Requires-Dist: trustedmcp>=0.0.4
-Requires-Dist: cookiecutter>=2.6.0
-Requires-Dist: omegaconf>=2.3.0
-Requires-Dist: llama-index>=0.14.0
-Requires-Dist: llama-index-embeddings-openai-like>=0.2.2
-Requires-Dist: llama-index-llms-openai-like>=0.5.1
-Requires-Dist: llama-index-vector-stores-opensearch>=0.6.1
-Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: agent-pilot-sdk==0.1.2
+Requires-Dist: fastmcp>=2.12.3
+Requires-Dist: trustedmcp==0.0.5
+Requires-Dist: cookiecutter==2.6.0
+Requires-Dist: omegaconf==2.3.0
+Requires-Dist: llama-index==0.14.0
+Requires-Dist: llama-index-embeddings-openai-like==0.2.2
+Requires-Dist: llama-index-llms-openai-like==0.5.1
+Requires-Dist: llama-index-vector-stores-opensearch==0.6.1
+Requires-Dist: psycopg2-binary==2.9.10
 Requires-Dist: asyncpg>=0.29.0
-Requires-Dist: pymysql>=1.1.1
-Requires-Dist: aiomysql>=0.3.2
+Requires-Dist: pymysql==1.1.1
+Requires-Dist: aiomysql==0.3.2
 Requires-Dist: opensearch-py==2.8.0
-Requires-Dist: filetype>=1.2.0
+Requires-Dist: filetype==1.2.0
 Requires-Dist: vikingdb-python-sdk>=0.1.3
-Requires-Dist: agentkit-sdk-python
+Requires-Dist: agentkit-sdk-python>=0.2.0
+Requires-Dist: python-frontmatter==1.1.0
 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: cozeloop>=0.1.21; extra == "extensions"
 Provides-Extra: database
 Requires-Dist: redis>=5.0; extra == "database"
 Requires-Dist: pymysql>=1.1.1; extra == "database"

{veadk_python-0.2.32 → veadk_python-0.5.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.32"
+version = "0.5.5"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -12,36 +12,35 @@ authors = [
     {name = "Meng Wang", email = "mengwangwm@gmail.com"}
 ]
 dependencies = [
-    "pydantic-settings>=2.10.1", # Config management
-    "a2a-sdk>=0.3.0", # For Google Agent2Agent protocol
-    "deprecated>=1.2.18",
-    "google-adk>=1.10.0", # For basic agent architecture
-    "google-adk<=1.19.0", # For basic agent architecture
+    "pydantic-settings==2.10.1", # Config management
+    "a2a-sdk>=0.3.7", # For Google Agent2Agent protocol
+    "deprecated==1.2.18",
+    "google-adk>=1.19.0", # For basic agent architecture
     "litellm>=1.74.3", # For model inference
-    "loguru>=0.7.3", # For better logging
-    "opentelemetry-exporter-otlp>=1.35.0",
+    "loguru==0.7.3", # For better logging
+    "opentelemetry-exporter-otlp==1.37.0",
     "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.33", # For Volcengine API
+    "wrapt==1.17.2", # For patching built-in functions
+    "volcengine-python-sdk>=5.0.1", # 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
-    "trustedmcp>=0.0.4", # For running TrustedMCP
-    "cookiecutter>=2.6.0", # For cloud deploy
-    "omegaconf>=2.3.0", # For agent builder
-    "llama-index>=0.14.0",
-    "llama-index-embeddings-openai-like>=0.2.2",
-    "llama-index-llms-openai-like>=0.5.1",
-    "llama-index-vector-stores-opensearch>=0.6.1",
-    "psycopg2-binary>=2.9.10", # For PostgreSQL database (short term memory)
+    "agent-pilot-sdk==0.1.2", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
+    "fastmcp>=2.12.3", # For running MCP
+    "trustedmcp==0.0.5", # For running TrustedMCP
+    "cookiecutter==2.6.0", # For cloud deploy
+    "omegaconf==2.3.0", # For agent builder
+    "llama-index==0.14.0",
+    "llama-index-embeddings-openai-like==0.2.2",
+    "llama-index-llms-openai-like==0.5.1",
+    "llama-index-vector-stores-opensearch==0.6.1",
+    "psycopg2-binary==2.9.10", # For PostgreSQL database (short term memory)
     "asyncpg>=0.29.0", # For async PostgreSQL database (short term memory)
-    "pymysql>=1.1.1", # For MySQL database (short term memory)
-    "aiomysql>=0.3.2", # For async MySQL database (short term memory)
+    "pymysql==1.1.1", # For MySQL database (short term memory)
+    "aiomysql==0.3.2", # For async MySQL database (short term memory)
     "opensearch-py==2.8.0",
-    "filetype>=1.2.0",
+    "filetype==1.2.0",
     "vikingdb-python-sdk>=0.1.3",
-    "agentkit-sdk-python"
+    "agentkit-sdk-python>=0.2.0",
+    "python-frontmatter==1.1.0",
 ]
 
 [project.scripts]
@@ -52,6 +51,7 @@ 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",
+    "cozeloop>=0.1.21", #For Cozeloop Prompt manager
 ]
 database = [
     "redis>=5.0",                 # For Redis database
@@ -74,7 +74,6 @@ dev = [
     "pytest-xdist>=3.8.0",
 ]
 
-
 [dependency-groups]
 dev = [
     "pre-commit>=4.2.0",            # Format checking

{veadk_python-0.2.32 → veadk_python-0.5.5}/veadk/agent.py

@@ -24,8 +24,11 @@ from typing import Optional, Union
 if not os.getenv("LITELLM_LOCAL_MODEL_COST_MAP"):
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
 
+import uuid
+
 from google.adk.agents import LlmAgent, RunConfig
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.agents.llm_agent import InstructionProvider, ToolUnion
 from google.adk.agents.run_config import StreamingMode
 from google.adk.models.lite_llm import LiteLlm
@@ -44,7 +47,11 @@ from veadk.knowledgebase import KnowledgeBase
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.processors import BaseRunProcessor, NoOpRunProcessor
-from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
+from veadk.prompts.agent_default_prompt import (
+    DEFAULT_DESCRIPTION,
+    DEFAULT_INSTRUCTION,
+)
+from veadk.prompts.prompt_manager import BasePromptManager
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_asyncio, patch_tracer
@@ -67,7 +74,7 @@ class Agent(LlmAgent):
         name (str): The name of the agent.
         description (str): A description of the agent, useful in A2A scenarios.
         instruction (Union[str, InstructionProvider]): The instruction or instruction provider.
-        model_name (str): Name of the model used by the agent.
+        model_name (Union[str, List[str]]): Name of the model used by the agent.
         model_provider (str): Provider of the model (e.g., openai).
         model_api_base (str): The base URL of the model API.
         model_api_key (str): The API key for accessing the model.
@@ -78,15 +85,20 @@ class Agent(LlmAgent):
         short_term_memory (Optional[ShortTermMemory]): Session-based memory for temporary context.
         long_term_memory (Optional[LongTermMemory]): Cross-session memory for persistent user context.
         tracers (list[BaseTracer]): List of tracers used for telemetry and monitoring.
+        enable_authz (bool): Whether to enable agent authorization checks.
+        auto_save_session (bool): Whether to automatically save sessions to long-term memory.
     """
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()).split("-")[0])
     name: str = DEFAULT_AGENT_NAME
     description: str = DEFAULT_DESCRIPTION
     instruction: Union[str, InstructionProvider] = DEFAULT_INSTRUCTION
 
-    model_name: str = Field(default_factory=lambda: settings.model.name)
+    model_name: Union[str, list[str]] = Field(
+        default_factory=lambda: settings.model.name
+    )
     model_provider: str = Field(default_factory=lambda: settings.model.provider)
     model_api_base: str = Field(default_factory=lambda: settings.model.api_base)
     model_api_key: str = Field(default_factory=lambda: settings.model.api_key)
@@ -96,6 +108,8 @@ class Agent(LlmAgent):
 
     sub_agents: list[BaseAgent] = Field(default_factory=list, exclude=True)
 
+    prompt_manager: Optional[BasePromptManager] = None
+
     knowledgebase: Optional[KnowledgeBase] = None
 
     short_term_memory: Optional[ShortTermMemory] = None
@@ -103,6 +117,10 @@ class Agent(LlmAgent):
 
     tracers: list[BaseTracer] = []
 
+    enable_responses: bool = False
+
+    context_cache_config: Optional[ContextCacheConfig] = None
+
     run_processor: Optional[BaseRunProcessor] = Field(default=None, exclude=True)
     """Optional run processor for intercepting and processing agent execution flows.
 
@@ -125,6 +143,10 @@ class Agent(LlmAgent):
 
     enable_authz: bool = False
 
+    auto_save_session: bool = False
+
+    skills: list[str] = Field(default_factory=list)
+
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
@@ -151,12 +173,41 @@ class Agent(LlmAgent):
         logger.info(f"Model extra config: {self.model_extra_config}")
 
         if not self.model:
-            self.model = LiteLlm(
-                model=f"{self.model_provider}/{self.model_name}",
-                api_key=self.model_api_key,
-                api_base=self.model_api_base,
-                **self.model_extra_config,
-            )
+            if self.enable_responses:
+                from veadk.models.ark_llm import ArkLlm
+
+                self.model = ArkLlm(
+                    model=f"{self.model_provider}/{self.model_name}",
+                    api_key=self.model_api_key,
+                    api_base=self.model_api_base,
+                    **self.model_extra_config,
+                )
+            else:
+                fallbacks = None
+                if isinstance(self.model_name, list):
+                    if self.model_name:
+                        model_name = self.model_name[0]
+                        fallbacks = [
+                            f"{self.model_provider}/{m}" for m in self.model_name[1:]
+                        ]
+                        logger.info(
+                            f"Using primary model: {model_name}, with fallbacks: {self.model_name[1:]}"
+                        )
+                    else:
+                        model_name = settings.model.name
+                        logger.warning(
+                            f"Empty model_name list provided, using default model from settings: {model_name}"
+                        )
+                else:
+                    model_name = self.model_name
+
+                self.model = LiteLlm(
+                    model=f"{self.model_provider}/{model_name}",
+                    api_key=self.model_api_key,
+                    api_base=self.model_api_base,
+                    fallbacks=fallbacks,
+                    **self.model_extra_config,
+                )
             logger.debug(
                 f"LiteLLM client created with config: {self.model_extra_config}"
             )
@@ -202,13 +253,70 @@ class Agent(LlmAgent):
             else:
                 self.before_agent_callback = check_agent_authorization
 
+        if self.prompt_manager:
+            self.instruction = self.prompt_manager.get_prompt
+
+        if self.auto_save_session:
+            if self.long_term_memory is None:
+                logger.warning(
+                    "auto_save_session is enabled, but long_term_memory is not initialized."
+                )
+            else:
+                from veadk.memory.save_session_callback import (
+                    save_session_to_long_term_memory,
+                )
+
+                if self.after_agent_callback:
+                    if isinstance(self.after_agent_callback, list):
+                        self.after_agent_callback.append(
+                            save_session_to_long_term_memory
+                        )
+                    else:
+                        self.after_agent_callback = [
+                            self.after_agent_callback,
+                            save_session_to_long_term_memory,
+                        ]
+                else:
+                    self.after_agent_callback = save_session_to_long_term_memory
+
+        if self.skills:
+            self.load_skills()
+
         logger.info(f"VeADK version: {VERSION}")
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
         logger.debug(
-            f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools'})}"
+            f"Agent: {self.model_dump(include={'id', 'name', 'model_name', 'model_api_base', 'tools', 'skills'})}"
+        )
+
+    def update_model(self, model_name: str):
+        logger.info(f"Updating model to {model_name}")
+        self.model = self.model.model_copy(
+            update={"model": f"{self.model_provider}/{model_name}"}
         )
 
+    def load_skills(self):
+        from pathlib import Path
+
+        from veadk.skills.utils import load_skills_from_directory
+
+        skills = []
+        for skill in self.skills:
+            path = Path(skill)
+            if path.is_dir():
+                skills.extend(load_skills_from_directory(path))
+            else:
+                logger.error(
+                    f"Skill {skill} is not a directory, skip. Loading skills from cloud is WIP."
+                )
+        if skills:
+            self.instruction += "\nYou have the following skills:\n"
+
+            for skill in skills:
+                self.instruction += (
+                    f"- name: {skill.name}\n- description: {skill.description}\n\n"
+                )
+
     async def _run(
         self,
         runner,
@@ -274,14 +382,20 @@ class Agent(LlmAgent):
             return
 
         if not self.tracers:
-            from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
+            from veadk.tracing.telemetry.opentelemetry_tracer import (
+                OpentelemetryTracer,
+            )
 
             self.tracers.append(OpentelemetryTracer())
 
         exporters = self.tracers[0].exporters  # type: ignore
 
-        from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
-        from veadk.tracing.telemetry.exporters.cozeloop_exporter import CozeloopExporter
+        from veadk.tracing.telemetry.exporters.apmplus_exporter import (
+            APMPlusExporter,
+        )
+        from veadk.tracing.telemetry.exporters.cozeloop_exporter import (
+            CozeloopExporter,
+        )
         from veadk.tracing.telemetry.exporters.tls_exporter import TLSExporter
 
         if enable_apmplus_tracer and not any(

veadk_python-0.5.5/veadk/auth/veauth/mse_veauth.py

@@ -0,0 +1,118 @@
+# 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 veadk.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.configs.database_configs import MSENacosConfig
+from veadk.utils.logger import get_logger
+from veadk.utils.volcengine_sign import ve_request
+
+logger = get_logger(__name__)
+
+
+def get_instance_id_by_name(instance_name: str, region: str, project_name: str) -> str:
+    logger.info(f"Fetching MSE Nacos instance id by instance name {instance_name} ...")
+
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_token = ""
+
+    if not (access_key and secret_key):
+        # try to get from vefaas iam
+        cred = get_credential_from_vefaas_iam()
+        access_key = cred.access_key_id
+        secret_key = cred.secret_access_key
+        session_token = cred.session_token
+
+    res = ve_request(
+        request_body={
+            "Filter": {"Status": [], "ProjectName": project_name},
+            "PageNumber": 1,
+            "PageSize": 10,
+            "ProjectName": project_name,
+        },
+        header={"X-Security-Token": session_token},
+        action="ListNacosRegistries",
+        ak=access_key,
+        sk=secret_key,
+        service="mse",
+        version="2022-01-01",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+
+    try:
+        for item in res["Result"]["Items"]:
+            if item["Name"] == instance_name:
+                logger.info(
+                    f"Found MSE Nacos instance id {item['Id']} by instance name {instance_name}"
+                )
+                return item["Id"]
+        raise ValueError(f"Id by instance name {instance_name} not found: {res}")
+    except Exception as e:
+        logger.error(
+            f"Failed to get MSE Nacos instance id by name {instance_name}: {e}, response: {res}"
+        )
+        raise e
+
+
+def get_mse_cridential(
+    instance_name: str,
+    region: str = "cn-beijing",
+    project_name: str = "default",
+) -> MSENacosConfig:
+    logger.info("Fetching MSE Nacos token...")
+
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_token = ""
+
+    if not (access_key and secret_key):
+        # try to get from vefaas iam
+        cred = get_credential_from_vefaas_iam()
+        access_key = cred.access_key_id
+        secret_key = cred.secret_access_key
+        session_token = cred.session_token
+
+    instance_id = get_instance_id_by_name(
+        instance_name=instance_name, region=region, project_name=project_name
+    )
+
+    res = ve_request(
+        request_body={
+            "Id": instance_id,
+            "ProjectName": project_name,
+        },
+        header={"X-Security-Token": session_token},
+        action="GetNacosRegistry",
+        ak=access_key,
+        sk=secret_key,
+        service="mse",
+        version="2022-01-01",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+
+    try:
+        logger.info(
+            f"Successfully fetched MSE Nacos endpoint {res['Result']['NacosRegistry']['PublicAddress']} and corresponding password."
+        )
+        return MSENacosConfig(
+            endpoint=res["Result"]["NacosRegistry"]["PublicAddress"],
+            password=res["Result"]["NacosRegistry"]["InitialPassword"],
+        )
+    except Exception as e:
+        logger.error(f"Failed to get MSE Nacos token: {e}, response: {res}")
+        raise e

{veadk_python-0.2.32 → veadk_python-0.5.5}/veadk/cli/cli_create.py

@@ -147,6 +147,10 @@ def create(agent_name: str, ark_api_key: str) -> None:
     """
     if not agent_name:
         agent_name = click.prompt("Enter the agent name")
+
+    if "-" in agent_name:
+        raise ValueError("Agent name cannot contain '-'. Use '_' instead.")
+
     if not ark_api_key:
         ark_api_key = _prompt_for_ark_api_key()
 

{veadk_python-0.2.32 → veadk_python-0.5.5}/veadk/cli/cli_rl.py

@@ -38,8 +38,8 @@ def rl_group():
     "--platform",
     "-p",
     required=True,
-    type=click.Choice(["ark"], case_sensitive=False),
-    help="Scaffold platform type (only support for now: ark)",
+    type=click.Choice(["ark", "lightning"], case_sensitive=False),
+    help="Scaffold platform type (supported: ark, lightning)",
 )
 @click.option(
     "--workspace", "-w", required=True, type=str, help="Target workspace directory name"
@@ -52,8 +52,9 @@ def rl_group():
 )
 def rl_init(platform: str, workspace: str, overwrite: bool):
     """
-    Initialize RL scaffold project for ark platform
+    Initialize RL scaffold project for ark or lightning platform
     Example: veadk rl init --platform ark --workspace veadk_rl_ark_project
+    Example: veadk rl init --platform lightning --workspace veadk_rl_lightning_project
     """
     # Locate template directory
     rl_template_root = get_rl_template_root()

{veadk_python-0.2.32 → veadk_python-0.5.5}/veadk/cli/cli_web.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from functools import wraps
 
 import click
@@ -139,5 +140,14 @@ def web(ctx, *args, **kwargs) -> None:
     # from Google ADK and Litellm
     if "--log_level" not in extra_args:
         extra_args.extend(["--log_level", "ERROR"])
+        logging.basicConfig(level=logging.ERROR, force=True)
+
+    if "--log_level" in extra_args:
+        logging.basicConfig(
+            level=getattr(
+                logging, extra_args[extra_args.index("--log_level") + 1].upper()
+            ),
+            force=True,
+        )
 
     cli_web.main(args=extra_args, standalone_mode=False)

{veadk_python-0.2.32 → veadk_python-0.5.5}/veadk/cli/templates/rl/ark/plugins/raw_async_veadk_rollout.py

@@ -15,7 +15,6 @@
 import json
 import asyncio
 from typing import Optional, Any, cast
-from loguru import logger
 from ark_sdk.resources.pipeline_plugin import rollout
 from ark_sdk.types.pipeline_plugin import PluginInstance, Runtime
 from ark_sdk.types.pipeline_plugin.pipeline_plugin import PluginContext
@@ -25,7 +24,6 @@ from ark_sdk.types.pipeline_plugin.rollout import (
     ChatCompletionResponse,
     RolloutInferenceProxy,
     RolloutResult,
-    PluginStatus,
 )
 from veadk.agent import Agent
 from veadk.memory.short_term_memory import ShortTermMemory
@@ -33,9 +31,9 @@ from veadk.runner import Runner
 from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
 from veadk.tracing.telemetry.exporters.cozeloop_exporter import CozeloopExporter
 from veadk.tracing.telemetry.exporters.cozeloop_exporter import CozeloopExporterConfig
+from veadk.tools.demo_tools import get_city_weather
 from google.adk.models.lite_llm import LiteLLMClient, LiteLlm
 from litellm import ModelResponse
-from cozeloop.decorator import observe
 
 # BASE_MODEL 格式 : "{model_provider}/{model_name}"
 BASE_MODEL = "openai/doubao-seed-1-6-flash-250615"
@@ -50,13 +48,6 @@ exporters = [CozeloopExporter(config=coze_config)]
 tracer = OpentelemetryTracer(exporters=cast(Any, exporters))
 
 
-@observe()
-def get_current_weather(location: str, unit="摄氏度"):
-    # 实际调用天气查询 API 的逻辑
-    # 此处为示例，返回模拟的天气数据
-    return f"{location}今天天气晴朗，温度 25 {unit}。"
-
-
 class RecordingLiteLlm(LiteLlm):
     """
     在调用 LiteLlm 的 completion/acompletion 时，拦截并记录原始 ModelResponse。
@@ -171,6 +162,7 @@ async def demo_veadk_rollout(
         model_provider="openai",
         model_api_key=proxy.jwt_token,
         tracers=[tracer],
+        tools=[get_city_weather],
         model=model_instance,
     )
 
@@ -224,29 +216,6 @@ async def demo_veadk_rollout(
         if model_response.choices[0].finish_reason != "tool_calls":
             # 模型最终总结，没有调用工具意愿
             break
-        tool_calls = model_response.choices[0].message.tool_calls
-        for tool_call in tool_calls or []:
-            tool_name = tool_call.function.name
-            if tool_name == "get_current_weather":
-                try:
-                    args = json.loads(tool_call.function.arguments)
-                    tool_result = get_current_weather(**args)
-                except Exception as e:
-                    logger.error(f"get_current_weather error: {e}")
-                    return RolloutResult(
-                        status=PluginStatus.SUCCESS,
-                        extra={
-                            "reward": -1,
-                        },
-                    )
-                # 将工具结果加入消息列表
-                messages.append(
-                    {
-                        "role": "tool",
-                        "content": tool_result,
-                        "tool_call_id": tool_call.id,
-                    }
-                )
     # 默认return None则视为rollout成功
     return None