veadk-python 0.2.11__tar.gz → 0.2.13__tar.gz

{veadk_python-0.2.11 → veadk_python-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.2.11
+Version: 0.2.13
 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

{veadk_python-0.2.11 → veadk_python-0.2.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.11"
+version = "0.2.13"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"

{veadk_python-0.2.11 → veadk_python-0.2.13}/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_python-0.2.13/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_python-0.2.11 → veadk_python-0.2.13}/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_python-0.2.11 → veadk_python-0.2.13}/veadk/integrations/ve_tos/ve_tos.py

@@ -36,11 +36,14 @@ class VeTOS:
         self,
         ak: str = "",
         sk: str = "",
+        session_token: str = "",
         region: str = "cn-beijing",
         bucket_name: str = DEFAULT_TOS_BUCKET_NAME,
     ) -> None:
         self.ak = ak if ak else os.getenv("VOLCENGINE_ACCESS_KEY", "")
         self.sk = sk if sk else os.getenv("VOLCENGINE_SECRET_KEY", "")
+        self.session_token = session_token
+
         # Add empty value validation
         if not self.ak or not self.sk:
             raise ValueError(
@@ -71,6 +74,7 @@ class VeTOS:
             self._client = self._tos_module.TosClientV2(
                 ak=self.ak,
                 sk=self.sk,
+                security_token=self.session_token,
                 endpoint=f"tos-{self.region}.volces.com",
                 region=self.region,
             )
@@ -85,6 +89,7 @@ class VeTOS:
             self._client = self._tos_module.TosClientV2(
                 self.ak,
                 self.sk,
+                security_token=self.session_token,
                 endpoint=f"tos-{self.region}.volces.com",
                 region=self.region,
             )

{veadk_python-0.2.11 → veadk_python-0.2.13}/veadk/knowledgebase/backends/utils.py

@@ -46,6 +46,7 @@ def build_vikingdb_knowledgebase_request(
     path: str,
     volcengine_access_key: str,
     volcengine_secret_key: str,
+    session_token: str = "",
     method: Literal["GET", "POST", "PUT", "DELETE"] = "POST",
     region: str = "cn-beijing",
     params=None,
@@ -85,7 +86,7 @@ def build_vikingdb_knowledgebase_request(
         r.set_body(json.dumps(data))
 
     credentials = Credentials(
-        volcengine_access_key, volcengine_secret_key, "air", region
+        volcengine_access_key, volcengine_secret_key, "air", region, session_token
     )
     SignerV4.sign(r, credentials)
     return r

{veadk_python-0.2.11 → veadk_python-0.2.13}/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import json
+import os
 import re
 from pathlib import Path
 from typing import Any, Literal
@@ -23,7 +24,7 @@ from pydantic import Field
 from typing_extensions import override
 
 import veadk.config  # noqa E401
-from veadk.config import getenv
+from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
 from veadk.configs.database_configs import NormalTOSConfig, TOSConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import build_vikingdb_knowledgebase_request
@@ -58,14 +59,16 @@ def get_files_in_directory(directory: str):
 
 
 class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
-    volcengine_access_key: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
+    volcengine_access_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
     )
 
-    volcengine_secret_key: str = Field(
-        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    volcengine_secret_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
     )
 
+    session_token: str = ""
+
     volcengine_project: str = "default"
     """VikingDB knowledgebase project in Volcengine console platform. Default by `default`"""
 
@@ -75,6 +78,15 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
     tos_config: TOSConfig | NormalTOSConfig = Field(default_factory=TOSConfig)
     """TOS config, used to upload files to TOS"""
 
+    def model_post_init(self, __context: Any) -> None:
+        self.precheck_index_naming()
+
+        # check whether collection exist, if not, create it
+        if not self.collection_status()["existed"]:
+            logger.warning(
+                f"VikingDB knowledgebase collection {self.index} does not exist, please create it first..."
+            )
+
     def precheck_index_naming(self):
         if not (
             isinstance(self.index, str)
@@ -86,18 +98,21 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
                 "it must start with an English letter, contain only letters, numbers, and underscores, and have a length of 1-128."
             )
 
-    def model_post_init(self, __context: Any) -> None:
-        self.precheck_index_naming()
-
-        # check whether collection exist, if not, create it
-        if not self.collection_status()["existed"]:
-            logger.warning(
-                f"VikingDB knowledgebase collection {self.index} does not exist, please create it first..."
-            )
-
-        self._tos_client = VeTOS(
-            ak=self.volcengine_access_key,
-            sk=self.volcengine_secret_key,
+    def _get_tos_client(self) -> VeTOS:
+        volcengine_access_key = self.volcengine_access_key
+        volcengine_secret_key = self.volcengine_secret_key
+        session_token = self.session_token
+
+        if not (volcengine_access_key and volcengine_secret_key):
+            cred = get_credential_from_vefaas_iam()
+            volcengine_access_key = cred.access_key_id
+            volcengine_secret_key = cred.secret_access_key
+            session_token = cred.session_token
+
+        return VeTOS(
+            ak=volcengine_access_key,
+            sk=volcengine_secret_key,
+            session_token=session_token,
             region=self.tos_config.region,
             bucket_name=self.tos_config.bucket,
         )
@@ -404,6 +419,8 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         metadata: dict | None = None,
     ) -> str:
         # Here, we set the metadata via the TOS object, ref: https://www.volcengine.com/docs/84313/1254624
+        self._tos_client = self._get_tos_client()
+
         self._tos_client.bucket_name = tos_bucket_name
         coro = self._tos_client.upload(
             object_key=object_key,
@@ -504,10 +521,21 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
     ) -> dict:
         VIKINGDB_KNOWLEDGEBASE_BASE_URL = "api-knowledgebase.mlp.cn-beijing.volces.com"
 
+        volcengine_access_key = self.volcengine_access_key
+        volcengine_secret_key = self.volcengine_secret_key
+        session_token = self.session_token
+
+        if not (volcengine_access_key and volcengine_secret_key):
+            cred = get_credential_from_vefaas_iam()
+            volcengine_access_key = cred.access_key_id
+            volcengine_secret_key = cred.secret_access_key
+            session_token = cred.session_token
+
         request = build_vikingdb_knowledgebase_request(
             path=path,
-            volcengine_access_key=self.volcengine_access_key,
-            volcengine_secret_key=self.volcengine_secret_key,
+            volcengine_access_key=volcengine_access_key,
+            volcengine_secret_key=volcengine_secret_key,
+            session_token=session_token,
             method=method,
             data=body,
         )

{veadk_python-0.2.11 → veadk_python-0.2.13}/veadk/memory/long_term_memory.py

@@ -104,6 +104,15 @@ class LongTermMemory(BaseMemoryService, BaseModel):
 
         # Once user define backend config, use it directly
         if self.backend_config:
+            if "index" not in self.backend_config:
+                logger.warning(
+                    "Attribute `index` not provided in backend_config, use `index` or `app_name` instead."
+                )
+                self.backend_config["index"] = self.index or self.app_name
+
+            logger.debug(
+                f"Init {self.backend}, Use provided backend config: {self.backend_config}"
+            )
             self._backend = _get_backend_cls(self.backend)(**self.backend_config)
             return
 

{veadk_python-0.2.11 → veadk_python-0.2.13}/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -49,7 +49,23 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
     region: str = "cn-beijing"
     """VikingDB memory region"""
 
+    memory_type: list[str] = Field(default_factory=list)
+
     def model_post_init(self, __context: Any) -> None:
+        # We get memory type from:
+        # 1. user input
+        # 2. environment variable
+        # 3. default value
+        if not self.memory_type:
+            env_memory_type = os.getenv("DATABASE_VIKINGMEM_MEMORY_TYPE")
+            if env_memory_type:
+                # "event_1, event_2" -> ["event_1", "event_2"]
+                self.memory_type = [x.strip() for x in env_memory_type.split(",")]
+            else:
+                self.memory_type = ["sys_event_v1", "event_v1"]
+
+        logger.info(f"Using memory type: {self.memory_type}")
+
         # check whether collection exist, if not, create it
         if not self._collection_exist():
             self._create_collection()
@@ -68,17 +84,23 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
         try:
             client = self._get_client()
             client.get_collection(collection_name=self.index)
+            logger.info(f"Collection {self.index} exist.")
             return True
         except Exception:
+            logger.info(f"Collection {self.index} not exist.")
             return False
 
     def _create_collection(self) -> None:
+        logger.info(
+            f"Create collection with collection_name={self.index}, builtin_event_types={self.memory_type}"
+        )
         client = self._get_client()
         response = client.create_collection(
             collection_name=self.index,
             description="Created by Volcengine Agent Development Kit VeADK",
-            builtin_event_types=["sys_event_v1"],
+            builtin_event_types=self.memory_type,
         )
+        logger.debug(f"Create collection with response {response}")
         return response
 
     def _get_client(self) -> VikingDBMemoryClient:
@@ -135,10 +157,7 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
     def search_memory(
         self, user_id: str, query: str, top_k: int, **kwargs
     ) -> list[str]:
-        filter = {
-            "user_id": user_id,
-            "memory_type": ["sys_event_v1"],
-        }
+        filter = {"user_id": user_id, "memory_type": self.memory_type}
 
         logger.debug(
             f"Request for search memory in VikingDB: filter={filter}, collection_name={self.index}, query={query}, limit={top_k}"