veadk-python 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

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.
 
@@ -196,7 +199,6 @@ class Agent(LlmAgent):
         collect_runtime_data: bool = False,
         eval_set_id: str = "",
         save_session_to_memory: bool = False,
-        enable_memory_optimization: bool = False,
     ):
         """Running the agent. The runner and session service will be created automatically.
 
@@ -226,7 +228,6 @@ class Agent(LlmAgent):
         # memory service
         short_term_memory = ShortTermMemory(
             backend="database" if load_history_sessions_from_db else "local",
-            enable_memory_optimization=enable_memory_optimization,
             db_url=db_url,
         )
         session_service = short_term_memory.session_service

veadk/auth/veauth/opensearch_veauth.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.
+
+# 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 OpensearchVeAuth(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 Opensearch STS 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:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token

veadk/auth/veauth/postgresql_veauth.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.
+
+# 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 PostgreSqlVeAuth(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 PostgreSQL STS 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:
+        if self._token:
+            return self._token
+        self._fetch_token()
+        return self._token

veadk/cli/cli.py

@@ -16,10 +16,11 @@
 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_pipeline import pipeline
 from veadk.cli.cli_prompt import prompt
 from veadk.cli.cli_web import web
-from veadk.cli.cli_pipeline import pipeline
 from veadk.version import VERSION
 
 
@@ -37,6 +38,7 @@ veadk.add_command(init)
 veadk.add_command(prompt)
 veadk.add_command(web)
 veadk.add_command(pipeline)
+veadk.add_command(eval)
 
 if __name__ == "__main__":
     veadk()

veadk/cli/cli_eval.py

@@ -0,0 +1,160 @@
+# 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 click
+
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+@click.command()
+@click.option(
+    "--agent-dir",
+    default=".",
+    help="To-be-evaluated agent directory. Must export `root_agent` in `agent.py`",
+)
+@click.option(
+    "--agent-a2a-url",
+    default=None,
+    help="To-be-evaluated agent URL. The agent should be deployed as A2A mode.",
+)
+@click.option(
+    "--evalset-file",
+    required=True,
+    help="Google ADK formatted evalset file path",
+)
+@click.option(
+    "--evaluator",
+    type=click.Choice(["adk", "deepeval"], case_sensitive=False),
+    help="Evaluator type, choose `adk` or `deepeval`",
+)
+@click.option(
+    "--judge-model-name",
+    default="doubao-1-5-pro-256k-250115",
+    help="Judge model name, default is `doubao-1-5-pro-256k-250115`. Useless under `adk` evaluator.",
+)
+@click.option(
+    "--volcengine-access-key",
+    default=None,
+    help="Volcengine access key for using Volcengine models",
+)
+@click.option(
+    "--volcengine-secret-key",
+    default=None,
+    help="Volcengine secret key for using Volcengine models",
+)
+def eval(
+    agent_dir: str,
+    agent_a2a_url: str,
+    evalset_file: str,
+    evaluator: str,
+    judge_model_name: str,
+    volcengine_access_key: str,
+    volcengine_secret_key: str,
+) -> None:
+    import asyncio
+    import os
+    from pathlib import Path
+
+    from google.adk.cli.utils.agent_loader import AgentLoader
+
+    from veadk.a2a.remote_ve_agent import RemoteVeAgent
+    from veadk.config import getenv, settings
+    from veadk.prompts.prompt_evaluator import eval_principle_prompt
+
+    try:
+        from deepeval.metrics import GEval, ToolCorrectnessMetric
+        from deepeval.test_case import LLMTestCaseParams
+
+        from veadk.evaluation.adk_evaluator import ADKEvaluator
+        from veadk.evaluation.deepeval_evaluator import DeepevalEvaluator
+    except ImportError:
+        raise ImportError(
+            "Please install veadk with `[evaluation]` extras, e.g., `pip install veadk-python[eval]`"
+        )
+
+    # ====== prepare agent instance ======
+    if not agent_dir and not agent_a2a_url:
+        raise ValueError(
+            "Option `--agent-dir` or  `--agent-a2a-url` should be provided one of them."
+        )
+
+    if agent_dir and agent_a2a_url:
+        logger.warning(
+            "`--agent-dir` and `--agent-a2a-url` are both provided, will use `--agent-a2a-url`."
+        )
+        agent_instance = RemoteVeAgent(name="a2a_agent", url=agent_a2a_url)
+        logger.info(f"Loaded agent from {agent_a2a_url}")
+
+    if not agent_dir and agent_a2a_url:
+        agent_instance = RemoteVeAgent(name="a2a_agent", url=agent_a2a_url)
+        logger.info(f"Loaded agent from {agent_a2a_url}")
+
+    if agent_dir and not agent_a2a_url:
+        agent_instance = AgentLoader(str(Path(agent_dir).parent.resolve())).load_agent(
+            str(Path(agent_dir).name)
+        )
+        logger.info(f"Loaded agent from {agent_dir}, agent name: {agent_instance.name}")
+
+    # ====== prepare envs ======
+    if volcengine_access_key and "VOLCENGINE_ACCESS_KEY" not in os.environ:
+        os.environ["VOLCENGINE_ACCESS_KEY"] = volcengine_access_key
+    if volcengine_secret_key and "VOLCENGINE_SECRET_KEY" not in os.environ:
+        os.environ["VOLCENGINE_SECRET_KEY"] = volcengine_secret_key
+
+    # ====== prepare evaluator instance ======
+    evaluator_instance = None
+    if evaluator == "adk" and judge_model_name:
+        logger.warning(
+            "Using Google ADK evaluator, `--judge-model-name` will be ignored."
+        )
+        evaluator_instance = ADKEvaluator(agent=agent_instance)
+
+        asyncio.run(evaluator_instance.evaluate(eval_set_file_path=evalset_file))
+
+    if evaluator == "deepeval":
+        if not volcengine_access_key:
+            volcengine_access_key = getenv("VOLCENGINE_ACCESS_KEY")
+        if not volcengine_secret_key:
+            volcengine_secret_key = getenv("VOLCENGINE_SECRET_KEY")
+
+        evaluator_instance = DeepevalEvaluator(
+            agent=agent_instance,
+            judge_model_api_key=settings.model.api_key,
+            judge_model_name=judge_model_name,
+        )
+
+        judge_model = evaluator_instance.judge_model
+
+        metrics = [
+            GEval(
+                threshold=0.8,
+                name="Base Evaluation",
+                criteria=eval_principle_prompt,
+                evaluation_params=[
+                    LLMTestCaseParams.INPUT,
+                    LLMTestCaseParams.ACTUAL_OUTPUT,
+                    LLMTestCaseParams.EXPECTED_OUTPUT,
+                ],
+                model=judge_model,
+            ),
+            ToolCorrectnessMetric(threshold=0.5),
+        ]
+
+        asyncio.run(
+            evaluator_instance.evaluate(
+                eval_set_file_path=evalset_file, metrics=metrics
+            )
+        )

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/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,7 @@ def web(host: str) -> None:
         eval_sets_manager: Any,
         eval_set_results_manager: Any,
         agents_dir: str,
+        extra_plugins: Optional[list[str]] = None,
     ):
         self.agent_loader = agent_loader
         self.artifact_service = artifact_service
@@ -162,6 +164,7 @@ 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 []
 
         # parse VeADK memories
         short_term_memory, long_term_memory = _get_memory(module_path=agents_dir)
@@ -178,4 +181,6 @@ def web(host: str) -> None:
     agents_dir = os.getcwd()
     logger.info(f"Load agents from {agents_dir}")
 
-    cli_tools_click.cli_web.main(args=[agents_dir, "--host", host])
+    cli_tools_click.cli_web.main(
+        args=[agents_dir, "--host", host, "--log_level", "ERROR"]
+    )

veadk/configs/database_configs.py

@@ -32,6 +32,8 @@ class OpensearchConfig(BaseSettings):
 
     password: str = ""
 
+    secret_token: str = ""
+
 
 class MysqlConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_MYSQL_")
@@ -46,6 +48,25 @@ class MysqlConfig(BaseSettings):
 
     charset: str = "utf8"
 
+    secret_token: str = ""
+    """STS token for MySQL auth, not supported yet."""
+
+
+class PostgreSqlConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="DATABASE_POSTGRESQL_")
+
+    host: str = ""
+
+    port: int = 5432
+
+    user: str = ""
+
+    password: str = ""
+
+    database: str = ""
+
+    secret_token: str = ""
+
 
 class RedisConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="DATABASE_REDIS_")
@@ -58,6 +79,18 @@ class RedisConfig(BaseSettings):
 
     db: int = 0
 
+    secret_token: str = ""
+    """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_")
@@ -81,3 +114,13 @@ class TOSConfig(BaseSettings):
 
         VeTOS(bucket_name=_bucket).create_bucket()
         return _bucket
+
+
+class NormalTOSConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_")
+
+    endpoint: str = "tos-cn-beijing.volces.com"
+
+    region: str = "cn-beijing"
+
+    bucket: str

veadk/configs/model_configs.py

@@ -40,3 +40,35 @@ class ModelConfig(BaseSettings):
     @cached_property
     def api_key(self) -> str:
         return os.getenv("MODEL_AGENT_API_KEY") or ARKVeAuth().token
+
+
+class EmbeddingModelConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="MODEL_EMBEDDING_")
+
+    name: str = "doubao-embedding-text-240715"
+    """Model name for embedding."""
+
+    dim: int = 2560
+    """Embedding dim is different from different models."""
+
+    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/"
+    """The api base of the model for embedding."""
+
+    @cached_property
+    def api_key(self) -> str:
+        return os.getenv("MODEL_EMBEDDING_API_KEY") or ARKVeAuth().token
+
+
+class NormalEmbeddingModelConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="MODEL_EMBEDDING_")
+
+    name: str = "doubao-embedding-text-240715"
+    """Model name for embedding."""
+
+    dim: int = 2560
+    """Embedding dim is different from different models."""
+
+    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/"
+    """The api base of the model for embedding."""
+
+    api_key: str

veadk/consts.py

@@ -27,14 +27,16 @@ 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": {
             "type": getenv("MODEL_AGENT_CACHING", "enabled"),
         },
-        "thinking": {
-            "type": "disabled"
-        },  # disable thinking for token saving and shorter TTFT by default
+        # "thinking": {
+        #     "type": "disabled"
+        # },
         "expire_at": int(time.time()) + 3600,  # expire after 1 hour
     },
 }
@@ -56,6 +58,11 @@ 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_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/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 = []