Merge branch 'volcengine:main' into main

Author: doraemonlove

.gitleaks.toml

@@ -73,4 +73,4 @@ description = "Empty environment variables with KEY pattern"
 regex = '''os\.environ\[".*?KEY"\]\s*=\s*".+"'''
 
 [allowlist]
-paths = ["requirements.txt"]
+paths = ["requirements.txt", "tests"]

docs/content/6.memory/3.long-term-memory.md

@@ -7,13 +7,13 @@ navigation:
 
 ## 使用方法
 
-VeADK 的长期记忆通常存储在数据库中，通过如下方式定义一个长期记忆：
+VeADK 的长期记忆通常存储在数据库中，你需要在初始化长期记忆时定义 `index` 来指定后端索引名称。通过如下方式定义一个长期记忆：
 
 ```python
 from veadk.memory.long_term_memory import LongTermMemory
 
-# 由于长期记忆需要构建索引，因此你必须在初始化长期记忆时定义 `app_name` 以及 `user_id`
-long_term_memory = LongTermMemory(app_name="my_app_name", user_id="user_id")
+# 
+long_term_memory = LongTermMemory(index="my_index")
 ```
 
 通过如下例子说明长期记忆：
@@ -32,7 +32,7 @@ user_id = "temp_user"
 teaching_session_id = "teaching_session"
 student_session_id = "student_session"
 
-long_term_memory = LongTermMemory(backend="local", app_name=app_name, user_id=user_id)
+long_term_memory = LongTermMemory(backend="local", index=app_name)
 
 agent = Agent(long_term_memory=long_term_memory)
 
@@ -90,8 +90,4 @@ print(response)
   ::field{name="app_name" type="string"}
   Agent 应用名称，用于多应用区分。默认空字符串。
   ::
-
-  ::field{name="user_id" type="string"}
-  Agent 用户 ID，用于区分不同用户的长期记忆。默认空字符串。
-  ::
 ::

docs/content/7.knowledgebase/1.knowledgebase.md

@@ -9,7 +9,7 @@ navigation:
 
 VeADK 基于 Llama-index 作为知识库的主要处理入口。开发者可上传文本、文件、目录，我们会为您进行自动切片。
 
-创建知识库时，您必须要提供您的 `app_name`（将会用来自动构建索引名称），或指定一个知识库的索引：
+创建知识库时，您必须要提供您的知识库后端索引名称 `index`，或指定 `app_name` 来作为索引名称：
 
 ```python
 from veadk.knowledgebase import KnowledgeBase

tests/test_agent.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from unittest.mock import Mock, patch
 
 from google.adk.agents.llm_agent import LlmAgent
@@ -33,11 +34,10 @@
 
 
 def test_agent():
-    knowledgebase = KnowledgeBase(
-        index="test_index",
-        backend="local",
-        backend_config={"embedding_config": {"api_key": "test"}},
-    )
+    os.environ["MODEL_EMBEDDING_API_KEY"] = "mocked_api_key"
+
+    knowledgebase = KnowledgeBase(index="test_index", backend="local")
+
     long_term_memory = LongTermMemory(backend="local")
     tracer = OpentelemetryTracer()
 
@@ -69,8 +69,6 @@ def test_agent():
 
     assert agent.knowledgebase == knowledgebase
     assert agent.knowledgebase.backend == "local"  # type: ignore
-    assert load_knowledgebase_tool.knowledgebase == agent.knowledgebase
-    assert load_knowledgebase_tool.load_knowledgebase_tool in agent.tools
 
     assert agent.long_term_memory.backend == "local"  # type: ignore
     assert load_memory in agent.tools

tests/test_knowledgebase.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 
 import pytest
 
@@ -21,11 +22,9 @@
 
 @pytest.mark.asyncio
 async def test_knowledgebase():
+    os.environ["MODEL_EMBEDDING_API_KEY"] = "mocked_api_key"
+
     app_name = "kb_test_app"
-    kb = KnowledgeBase(
-        backend="local",
-        app_name=app_name,
-        backend_config={"embedding_config": {"api_key": "test"}},
-    )
+    kb = KnowledgeBase(backend="local", app_name=app_name)
 
     assert isinstance(kb._backend, InMemoryKnowledgeBackend)

tests/test_long_term_memory.py

@@ -13,23 +13,20 @@
 # limitations under the License.
 
 
+import os
+
 import pytest
 from google.adk.tools import load_memory
 
 from veadk.agent import Agent
 from veadk.memory.long_term_memory import LongTermMemory
 
-app_name = "test_ltm"
-user_id = "test_user"
-
 
 @pytest.mark.asyncio
 async def test_long_term_memory():
-    long_term_memory = LongTermMemory(
-        backend="local",
-        # app_name=app_name,
-        # user_id=user_id,
-    )
+    os.environ["MODEL_EMBEDDING_API_KEY"] = "mocked_api_key"
+    long_term_memory = LongTermMemory(backend="local")
+
     agent = Agent(
         name="all_name",
         model_name="test_model_name",
@@ -43,7 +40,8 @@ async def test_long_term_memory():
 
     assert load_memory in agent.tools, "load_memory tool not found in agent tools"
 
-    assert not agent.long_term_memory._backend
+    assert agent.long_term_memory
+    assert agent.long_term_memory._backend
 
     # assert agent.long_term_memory._backend.index == build_long_term_memory_index(
     #     app_name, user_id

tests/test_runner.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 from google.genai import types
 
 from veadk.agent import Agent
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
-from veadk.runner import Runner
-
 
 # Import the standalone function instead of accessing as class method
-from veadk.runner import _convert_messages
+from veadk.runner import Runner, _convert_messages
 
 
 def _test_convert_messages(runner):
@@ -67,6 +67,8 @@ def _test_convert_messages(runner):
 
 def test_runner():
     """Test Runner class initialization and core properties"""
+    os.environ["MODEL_EMBEDDING_API_KEY"] = "mocked_api_key"
+
     short_term_memory = ShortTermMemory()
     long_term_memory = LongTermMemory(backend="local")
     agent = Agent(

veadk/agent.py

@@ -133,10 +133,14 @@ def model_post_init(self, __context: Any) -> None:
             )
 
         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,
+        )