veadk-python 0.2.10__py3-none-any.whl → 0.2.11__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/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)
 

veadk/memory/long_term_memory.py

@@ -72,10 +72,6 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
     raise ValueError(f"Unsupported long term memory backend: {backend}")
 
 
-def build_long_term_memory_index(app_name: str, user_id: str):
-    return f"{app_name}_{user_id}"
-
-
 class LongTermMemory(BaseMemoryService, BaseModel):
     backend: Union[
         Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
@@ -89,54 +85,48 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     top_k: int = 5
     """Number of top similar documents to retrieve during search."""
 
+    index: str = ""
+
     app_name: str = ""
 
     user_id: str = ""
+    """Deprecated attribute"""
 
     def model_post_init(self, __context: Any) -> None:
-        if self.backend == "viking_mem":
-            logger.warning(
-                "The `viking_mem` backend is deprecated, please use `viking` instead."
-            )
-            self.backend = "viking"
-
-        self._backend = None
-
         # Once user define a backend instance, use it directly
         if isinstance(self.backend, BaseLongTermMemoryBackend):
             self._backend = self.backend
+            self.index = self._backend.index
             logger.info(
-                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}"
+                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}"
             )
             return
 
+        # Once user define backend config, use it directly
         if self.backend_config:
-            logger.warning(
-                f"Initialized long term memory backend {self.backend} with config. We will ignore `app_name` and `user_id` if provided."
-            )
             self._backend = _get_backend_cls(self.backend)(**self.backend_config)
-            _index = self.backend_config.get("index", None)
-            if _index:
-                self._index = _index
-                logger.info(f"Long term memory index set to {self._index}.")
-            else:
-                logger.warning(
-                    "Cannot find index via backend_config, please set `index` parameter."
-                )
             return
 
-        if self.app_name and self.user_id:
-            self._index = build_long_term_memory_index(
-                app_name=self.app_name, user_id=self.user_id
-            )
-            logger.info(f"Long term memory index set to {self._index}.")
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
+        # Check index
+        self.index = self.index or self.app_name
+        if not self.index:
+            logger.warning(
+                "Attribute `index` or `app_name` not provided, use `default_app` instead."
             )
-        else:
+            self.index = "default_app"
+
+        # Forward compliance
+        if self.backend == "viking_mem":
             logger.warning(
-                "Neither `backend_instance`, `backend_config`, nor (`app_name`/`user_id`) is provided, the long term memory storage will initialize when adding a session."
+                "The `viking_mem` backend is deprecated, change to `viking` instead."
             )
+            self.backend = "viking"
+
+        self._backend = _get_backend_cls(self.backend)(index=self.index)
+
+        logger.info(
+            f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}"
+        )
 
     def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
         final_events = []
@@ -164,75 +154,32 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         self,
         session: Session,
     ):
-        app_name = session.app_name
         user_id = session.user_id
-
-        if not self._backend and isinstance(self.backend, str):
-            self._index = build_long_term_memory_index(app_name, user_id)
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
-            )
-            logger.info(
-                f"Initialize long term memory backend now, index is {self._index}"
-            )
-
-        if not self._index and self._index != build_long_term_memory_index(
-            app_name, user_id
-        ):
-            logger.warning(
-                f"The `app_name` or `user_id` is different from the initialized one, skip add session to memory. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}"
-            )
-            return
         event_strings = self._filter_and_convert_events(session.events)
 
         logger.info(
-            f"Adding {len(event_strings)} events to long term memory: index={self._index}"
+            f"Adding {len(event_strings)} events to long term memory: index={self.index}"
+        )
+        self._backend.save_memory(user_id=user_id, event_strings=event_strings)
+        logger.info(
+            f"Added {len(event_strings)} events to long term memory: index={self.index}, user_id={user_id}"
         )
-
-        if self._backend:
-            self._backend.save_memory(event_strings=event_strings, user_id=user_id)
-
-            logger.info(
-                f"Added {len(event_strings)} events to long term memory: index={self._index}"
-            )
-        else:
-            logger.error(
-                "Long term memory backend initialize failed, cannot add session to memory."
-            )
 
     @override
-    async def search_memory(self, *, app_name: str, user_id: str, query: str):
-        # prevent model invoke `load_memory` before add session to this memory
-        if not self._backend and isinstance(self.backend, str):
-            self._index = build_long_term_memory_index(app_name, user_id)
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
-            )
-            logger.info(
-                f"Initialize long term memory backend now, index is {self._index}"
-            )
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> SearchMemoryResponse:
+        logger.info(f"Search memory with query={query}")
 
-        if not self._index and self._index != build_long_term_memory_index(
-            app_name, user_id
-        ):
-            logger.warning(
-                f"The `app_name` or `user_id` is different from the initialized one. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}. Search memory return empty list."
+        memory_chunks = []
+        try:
+            memory_chunks = self._backend.search_memory(
+                query=query, top_k=self.top_k, user_id=user_id
             )
-            return SearchMemoryResponse(memories=[])
-
-        if not self._backend:
+        except Exception as e:
             logger.error(
-                "Long term memory backend is not initialized, cannot search memory."
+                f"Exception orrcus during memory search: {e}. Return empty memory chunks"
             )
-            return SearchMemoryResponse(memories=[])
-
-        logger.info(
-            f"Searching long term memory: query={query} index={self._index} top_k={self.top_k}"
-        )
-
-        memory_chunks = self._backend.search_memory(
-            query=query, top_k=self.top_k, user_id=user_id
-        )
 
         memory_events = []
         for memory in memory_chunks:
@@ -260,6 +207,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             )
 
         logger.info(
-            f"Return {len(memory_events)} memory events for query: {query} index={self._index}"
+            f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}"
         )
         return SearchMemoryResponse(memories=memory_events)

veadk/memory/long_term_memory_backends/base_backend.py

@@ -25,9 +25,11 @@ class BaseLongTermMemoryBackend(ABC, BaseModel):
         """Check the index name is valid or not"""
 
     @abstractmethod
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         """Save memory to long term memory backend"""
 
     @abstractmethod
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         """Retrieve memory from long term memory backend"""

veadk/memory/long_term_memory_backends/in_memory_backend.py

@@ -29,10 +29,6 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
-        # no checking
-        pass
-
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(
             model_name=self.embedding_config.name,
@@ -41,8 +37,12 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         )
         self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
 
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         for event_string in event_strings:
             document = Document(text=event_string)
             nodes = self._split_documents([document])
@@ -50,7 +50,9 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         return True
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
         retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 from typing import Any
-from typing_extensions import override
+
 from pydantic import Field
+from typing_extensions import override
 
 from veadk.configs.database_configs import Mem0Config
-
-
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
@@ -66,7 +65,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         pass
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(
+        self, event_strings: list[str], user_id: str = "default_user", **kwargs
+    ) -> bool:
         """Save memory to Mem0
 
         Args:
@@ -76,8 +77,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         Returns:
             bool: True if saved successfully, False otherwise
         """
-        user_id = kwargs.get("user_id", "default_user")
-
         try:
             logger.info(
                 f"Saving {len(event_strings)} events to Mem0 for user: {user_id}"
@@ -100,7 +99,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
             return False
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, query: str, top_k: int, user_id: str = "default_user", **kwargs
+    ) -> list[str]:
         """Search memory from Mem0
 
         Args:
@@ -111,7 +112,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         Returns:
             list[str]: List of memory strings
         """
-        user_id = kwargs.get("user_id", "default_user")
 
         try:
             logger.info(