veadk-python 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

veadk/agent.py

@@ -133,10 +133,14 @@ class Agent(LlmAgent):
             )
 
         if self.knowledgebase:
-            from veadk.tools import load_knowledgebase_tool
+            from veadk.tools.builtin_tools.load_knowledgebase import (
+                LoadKnowledgebaseTool,
+            )
 
-            load_knowledgebase_tool.knowledgebase = self.knowledgebase
-            self.tools.append(load_knowledgebase_tool.load_knowledgebase_tool)
+            load_knowledgebase_tool = LoadKnowledgebaseTool(
+                knowledgebase=self.knowledgebase
+            )
+            self.tools.append(load_knowledgebase_tool)
 
         if self.long_term_memory is not None:
             from google.adk.tools import load_memory

veadk/auth/veauth/ark_veauth.py

@@ -14,64 +14,56 @@
 
 import os
 
-from typing_extensions import override
-
-from veadk.auth.veauth.base_veauth import BaseVeAuth
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 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)
+def get_ark_token(region: str = "cn-beijing") -> str:
+    logger.info("Fetching ARK token...")
 
-        self._token: str = ""
+    access_key = os.getenv("VOLCENGINE_ACCESS_KEY")
+    secret_key = os.getenv("VOLCENGINE_SECRET_KEY")
+    session_token = ""
 
-    @override
-    def _fetch_token(self) -> None:
-        logger.info("Fetching ARK token...")
-        # 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}")
+    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
 
-        # 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}")
+    res = ve_request(
+        request_body={"ProjectName": "default", "Filter": {}},
+        header={"X-Security-Token": session_token},
+        action="ListApiKeys",
+        ak=access_key,
+        sk=secret_key,
+        service="ark",
+        version="2024-01-01",
+        region=region,
+        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}")
 
-    @property
-    def token(self) -> str:
-        if self._token:
-            return self._token
-        self._fetch_token()
-        return self._token
+    # get raw api key
+    res = ve_request(
+        request_body={"Id": first_api_key_id},
+        header={"X-Security-Token": session_token},
+        action="GetRawApiKey",
+        ak=access_key,
+        sk=secret_key,
+        service="ark",
+        version="2024-01-01",
+        region=region,
+        host="open.volcengineapi.com",
+    )
+    try:
+        return res["Result"]["ApiKey"]
+    except KeyError:
+        raise ValueError(f"Failed to get ARK api key: {res}")

veadk/auth/veauth/utils.py

@@ -0,0 +1,57 @@
+# 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 pathlib import Path
+
+from pydantic import BaseModel
+
+from veadk.consts import VEFAAS_IAM_CRIDENTIAL_PATH
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VeIAMCredential(BaseModel):
+    access_key_id: str
+    secret_access_key: str
+    session_token: str
+
+
+def get_credential_from_vefaas_iam() -> VeIAMCredential:
+    """Get credential from VeFaaS IAM file"""
+    logger.info(
+        f"Get Volcegnine access key or secret key from environment variables failed, try to get from VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH})."
+    )
+
+    path = Path(VEFAAS_IAM_CRIDENTIAL_PATH)
+
+    if not path.exists():
+        logger.error(
+            f"Get Volcegnine access key or secret key from environment variables failed, and VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH}) not exists. Please check your configuration."
+        )
+        raise FileNotFoundError(
+            f"Get Volcegnine access key or secret key from environment variables failed, and VeFaaS IAM file (path={VEFAAS_IAM_CRIDENTIAL_PATH}) not exists. Please check your configuration."
+        )
+
+    with open(VEFAAS_IAM_CRIDENTIAL_PATH, "r") as f:
+        cred_dict = json.load(f)
+        access_key = cred_dict["access_key_id"]
+        secret_key = cred_dict["secret_access_key"]
+        session_token = cred_dict["session_token"]
+        return VeIAMCredential(
+            access_key_id=access_key,
+            secret_access_key=secret_key,
+            session_token=session_token,
+        )

veadk/cli/cli.py

@@ -22,6 +22,7 @@ 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
+from veadk.cli.cli_uploadevalset import uploadevalset
 from veadk.version import VERSION
 
 
@@ -41,6 +42,7 @@ veadk.add_command(web)
 veadk.add_command(pipeline)
 veadk.add_command(eval)
 veadk.add_command(kb)
+veadk.add_command(uploadevalset)
 
 if __name__ == "__main__":
     veadk()

veadk/cli/cli_uploadevalset.py

@@ -0,0 +1,125 @@
+# 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
+import json
+import requests
+from veadk.utils.logger import get_logger
+from veadk.config import getenv
+from pathlib import Path
+
+logger = get_logger(__name__)
+
+
+@click.command()
+@click.option("--file", required=True, help="JSON file path containing dataset items")
+@click.option("--cozeloop-workspace-id", default=None, help="CozeLoop workspace ID")
+@click.option("--cozeloop-evalset-id", default=None, help="CozeLoop evaluation set ID")
+@click.option(
+    "--cozeloop-api-key",
+    default=None,
+    help="CozeLoop API key (or set COZELOOP_API_KEY env var)",
+)
+def uploadevalset(
+    file: str,
+    cozeloop_workspace_id: str,
+    cozeloop_evalset_id: str,
+    cozeloop_api_key: str,
+) -> None:
+    """Upload dataset items to CozeLoop evaluation set."""
+
+    if not cozeloop_workspace_id:
+        cozeloop_workspace_id = getenv(
+            "OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME"
+        )
+    if not cozeloop_evalset_id:
+        cozeloop_evalset_id = getenv("OBSERVABILITY_OPENTELEMETRY_COZELOOP_EVALSET_ID")
+    if not cozeloop_api_key:
+        cozeloop_api_key = getenv("OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY")
+
+    # Read JSON file
+    file_path = Path(file)
+    if not file_path.exists():
+        logger.error(f"File not found: {file}")
+        return
+
+    logger.info(f"Reading dataset from {file}")
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    # Prepare items
+    items = []
+    for case in data.get("eval_cases", []):
+        conversation = case.get("conversation", [])
+        for turn in conversation:
+            user_text = (
+                turn.get("user_content", {}).get("parts", [{}])[0].get("text", "")
+            )
+            output_text = (
+                turn.get("final_response", {}).get("parts", [{}])[0].get("text", "")
+            )
+
+            items.append(
+                {
+                    "turns": [
+                        {
+                            "field_datas": [
+                                {
+                                    "name": "input",
+                                    "content": {
+                                        "content_type": "Text",
+                                        "text": user_text,
+                                    },
+                                },
+                                {
+                                    "name": "output",
+                                    "content": {
+                                        "content_type": "Text",
+                                        "text": output_text,
+                                    },
+                                },
+                            ]
+                        }
+                    ]
+                }
+            )
+
+    # Upload to CozeLoop
+    url = f"https://api.coze.cn/v1/loop/evaluation/evaluation_sets/{cozeloop_evalset_id}/items"
+    logger.info(
+        f"Uploading {len(items)} items to workspace_id={cozeloop_workspace_id} evalset_id={cozeloop_evalset_id}"
+    )
+
+    response = requests.post(
+        url=url,
+        headers={
+            "Authorization": f"Bearer {cozeloop_api_key}",
+            "Content-Type": "application/json",
+            "X-TT-ENV": "ppe_eval_openapi",
+            "x-use-ppe": "1",
+        },
+        json={
+            "workspace_id": cozeloop_workspace_id,
+            "is_allow_partial_add": True,
+            "is_skip_invalid_items": True,
+            "items": items,
+        },
+    )
+
+    if response.status_code == 200:
+        logger.info(
+            f"Successfully uploaded dataset to CozeLoop evalset {cozeloop_evalset_id}"
+        )
+    else:
+        logger.error(f"Failed to upload dataset: {response.text}")

veadk/cli/cli_web.py

@@ -56,10 +56,14 @@ def _get_ltm_from_env() -> LongTermMemory | None:
     logger = get_logger(__name__)
 
     long_term_memory_backend = os.getenv("LONG_TERM_MEMORY_BACKEND")
+    app_name = os.getenv("VEADK_WEB_APP_NAME", "")
+    user_id = os.getenv("VEADK_WEB_USER_ID", "")
 
     if long_term_memory_backend:
         logger.info(f"Long term memory: backend={long_term_memory_backend}")
-        return LongTermMemory(backend=long_term_memory_backend)  # type: ignore
+        return LongTermMemory(
+            backend=long_term_memory_backend, app_name=app_name, user_id=user_id
+        )  # type: ignore
     else:
         logger.warning("No long term memory backend settings detected.")
         return None
@@ -131,7 +135,13 @@ def patch_adkwebserver_disable_openapi():
 
 @click.command()
 @click.option("--host", default="127.0.0.1", help="Host to run the web server on")
-def web(host: str) -> None:
+@click.option(
+    "--app_name", default="", help="The `app_name` for initializing long term memory"
+)
+@click.option(
+    "--user_id", default="", help="The `user_id` for initializing long term memory"
+)
+def web(host: str, app_name: str, user_id: str) -> None:
     """Launch web with long term and short term memory."""
     import os
     from typing import Any
@@ -175,6 +185,9 @@ def web(host: str) -> None:
         self.session_service = short_term_memory.session_service
         self.memory_service = long_term_memory
 
+    os.environ["VEADK_WEB_APP_NAME"] = app_name
+    os.environ["VEADK_WEB_USER_ID"] = user_id
+
     import google.adk.cli.adk_web_server
 
     google.adk.cli.adk_web_server.AdkWebServer.__init__ = init_for_veadk

veadk/configs/model_configs.py

@@ -17,7 +17,7 @@ from functools import cached_property
 
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from veadk.auth.veauth.ark_veauth import ARKVeAuth
+from veadk.auth.veauth.ark_veauth import get_ark_token
 from veadk.consts import (
     DEFAULT_MODEL_AGENT_API_BASE,
     DEFAULT_MODEL_AGENT_NAME,
@@ -39,7 +39,7 @@ class ModelConfig(BaseSettings):
 
     @cached_property
     def api_key(self) -> str:
-        return os.getenv("MODEL_AGENT_API_KEY") or ARKVeAuth().token
+        return os.getenv("MODEL_AGENT_API_KEY") or get_ark_token()
 
 
 class EmbeddingModelConfig(BaseSettings):
@@ -56,7 +56,7 @@ class EmbeddingModelConfig(BaseSettings):
 
     @cached_property
     def api_key(self) -> str:
-        return os.getenv("MODEL_EMBEDDING_API_KEY") or ARKVeAuth().token
+        return os.getenv("MODEL_EMBEDDING_API_KEY") or get_ark_token()
 
 
 class NormalEmbeddingModelConfig(BaseSettings):

veadk/consts.py

@@ -63,6 +63,15 @@ 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_TEXT_TO_IMAGE_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_IMAGE_EDIT_MODEL_NAME = "doubao-seededit-3-0-i2i-250628"
+DEFAULT_IMAGE_EDIT_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_VIDEO_MODEL_NAME = "doubao-seedance-1-0-pro-250528"
+DEFAULT_VIDEO_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
 DEFAULT_IMAGE_GENERATE_MODEL_NAME = "doubao-seedream-4-0-250828"
+DEFAULT_IMAGE_GENERATE_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+
+VEFAAS_IAM_CRIDENTIAL_PATH = "/var/run/secrets/iam/credential"

veadk/knowledgebase/knowledgebase.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Literal
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, Union
 
 from pydantic import BaseModel, Field
-from typing_extensions import Union
 
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.entry import KnowledgebaseEntry
@@ -54,11 +55,11 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
     raise ValueError(f"Unsupported knowledgebase backend: {backend}")
 
 
-def build_knowledgebase_index(app_name: str):
-    return f"veadk_kb_{app_name}"
+class KnowledgeBase(BaseModel):
+    name: str = "user_knowledgebase"
 
+    description: str = "This knowledgebase stores some user-related information."
 
-class KnowledgeBase(BaseModel):
     backend: Union[
         Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
     ] = "local"
@@ -73,9 +74,7 @@ class KnowledgeBase(BaseModel):
     """Configuration for the backend"""
 
     top_k: int = 10
-    """Number of top similar documents to retrieve during search.
-
-    Default is 10."""
+    """Number of top similar documents to retrieve during search"""
 
     app_name: str = ""
 
@@ -85,38 +84,27 @@ class KnowledgeBase(BaseModel):
     def model_post_init(self, __context: Any) -> None:
         if isinstance(self.backend, BaseKnowledgebaseBackend):
             self._backend = self.backend
+            self.index = self._backend.index
             logger.info(
                 f"Initialized knowledgebase with provided backend instance {self._backend.__class__.__name__}"
             )
             return
 
-        # must provide at least one of them
-        if not self.app_name and not self.index:
-            raise ValueError(
-                "Either `app_name` or `index` must be provided one of them."
-            )
-
-        # priority use index
-        if self.app_name and self.index:
-            logger.warning(
-                "`app_name` and `index` are both provided, using `index` as the knowledgebase index name."
-            )
+        # Once user define backend config, use it directly
+        if self.backend_config:
+            self._backend = _get_backend_cls(self.backend)(**self.backend_config)
+            return
 
-        # generate index name if `index` not provided but `app_name` is provided
-        if self.app_name and not self.index:
-            self.index = build_knowledgebase_index(self.app_name)
-            logger.info(
-                f"Knowledgebase index is set to {self.index} (generated by the app_name: {self.app_name})."
-            )
+        self.index = self.index or self.app_name
+        if not self.index:
+            raise ValueError("Either `index` or `app_name` must be provided.")
 
         logger.info(
-            f"Initializing knowledgebase: backend={self.backend} top_k={self.top_k}"
-        )
-        self._backend = _get_backend_cls(self.backend)(
-            index=self.index, **self.backend_config if self.backend_config else {}
+            f"Initializing knowledgebase: backend={self.backend} index={self.index} top_k={self.top_k}"
         )
+        self._backend = _get_backend_cls(self.backend)(index=self.index)
         logger.info(
-            f"Initialized knowledgebase with backend {self._backend.__class__.__name__}"
+            f"Initialized knowledgebase with backend {self.backend.__class__.__name__}"
         )
 
     def add_from_directory(self, directory: str, **kwargs) -> bool:
@@ -133,8 +121,7 @@ class KnowledgeBase(BaseModel):
 
     def search(self, query: str, top_k: int = 0, **kwargs) -> list[KnowledgebaseEntry]:
         """Search knowledge from knowledgebase"""
-        if top_k == 0:
-            top_k = self.top_k
+        top_k = top_k if top_k != 0 else self.top_k
 
         _entries = self._backend.search(query=query, top_k=top_k, **kwargs)