veadk-python 0.5.1__tar.gz → 0.5.2__tar.gz

{veadk_python-0.5.1/veadk_python.egg-info → veadk_python-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.5.1
+Version: 0.5.2
 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
@@ -211,18 +211,18 @@ License-File: LICENSE
 Requires-Dist: pydantic-settings==2.10.1
 Requires-Dist: a2a-sdk==0.3.7
 Requires-Dist: deprecated==1.2.18
-Requires-Dist: google-adk==1.19.0
-Requires-Dist: litellm==1.74.3
+Requires-Dist: google-adk>=1.19.0
+Requires-Dist: litellm>=1.74.3
 Requires-Dist: loguru==0.7.3
 Requires-Dist: opentelemetry-exporter-otlp==1.37.0
 Requires-Dist: opentelemetry-instrumentation-logging>=0.56b0
 Requires-Dist: wrapt==1.17.2
 Requires-Dist: openai<1.100
-Requires-Dist: volcengine-python-sdk==4.0.33
-Requires-Dist: volcengine==1.0.193
+Requires-Dist: volcengine-python-sdk>=5.0.1
+Requires-Dist: volcengine>=1.0.193
 Requires-Dist: agent-pilot-sdk==0.1.2
-Requires-Dist: fastmcp==2.12.3
-Requires-Dist: trustedmcp==0.0.4
+Requires-Dist: fastmcp>=2.12.3
+Requires-Dist: trustedmcp==0.0.5
 Requires-Dist: cookiecutter==2.6.0
 Requires-Dist: omegaconf==2.3.0
 Requires-Dist: llama-index==0.14.0
@@ -230,7 +230,7 @@ Requires-Dist: llama-index-embeddings-openai-like==0.2.2
 Requires-Dist: llama-index-llms-openai-like==0.5.1
 Requires-Dist: llama-index-vector-stores-opensearch==0.6.1
 Requires-Dist: psycopg2-binary==2.9.10
-Requires-Dist: asyncpg==0.29.0
+Requires-Dist: asyncpg>=0.29.0
 Requires-Dist: pymysql==1.1.1
 Requires-Dist: aiomysql==0.3.2
 Requires-Dist: opensearch-py==2.8.0
@@ -241,6 +241,7 @@ Provides-Extra: extensions
 Requires-Dist: redis>=5.0; extra == "extensions"
 Requires-Dist: tos>=2.8.4; extra == "extensions"
 Requires-Dist: llama-index-vector-stores-redis>=0.6.1; extra == "extensions"
+Requires-Dist: cozeloop>=0.1.21; extra == "extensions"
 Provides-Extra: database
 Requires-Dist: redis>=5.0; extra == "database"
 Requires-Dist: pymysql>=1.1.1; extra == "database"

{veadk_python-0.5.1 → veadk_python-0.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.5.1"
+version = "0.5.2"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -15,18 +15,18 @@ dependencies = [
     "pydantic-settings==2.10.1", # Config management
     "a2a-sdk==0.3.7", # For Google Agent2Agent protocol
     "deprecated==1.2.18",
-    "google-adk==1.19.0", # For basic agent architecture
-    "litellm==1.74.3", # For model inference
+    "google-adk>=1.19.0", # For basic agent architecture
+    "litellm>=1.74.3", # For model inference
     "loguru==0.7.3", # For better logging
     "opentelemetry-exporter-otlp==1.37.0",
     "opentelemetry-instrumentation-logging>=0.56b0",
     "wrapt==1.17.2", # For patching built-in functions
     "openai<1.100", # For fix https://github.com/BerriAI/litellm/issues/13710
-    "volcengine-python-sdk==4.0.33", # For Volcengine API
-    "volcengine==1.0.193", # For Volcengine sign
+    "volcengine-python-sdk>=5.0.1", # For Volcengine API
+    "volcengine>=1.0.193", # For Volcengine sign
     "agent-pilot-sdk==0.1.2", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
-    "fastmcp==2.12.3", # For running MCP
-    "trustedmcp==0.0.4", # For running TrustedMCP
+    "fastmcp>=2.12.3", # For running MCP
+    "trustedmcp==0.0.5", # For running TrustedMCP
     "cookiecutter==2.6.0", # For cloud deploy
     "omegaconf==2.3.0", # For agent builder
     "llama-index==0.14.0",
@@ -34,7 +34,7 @@ dependencies = [
     "llama-index-llms-openai-like==0.5.1",
     "llama-index-vector-stores-opensearch==0.6.1",
     "psycopg2-binary==2.9.10", # For PostgreSQL database (short term memory)
-    "asyncpg==0.29.0", # For async PostgreSQL database (short term memory)
+    "asyncpg>=0.29.0", # For async PostgreSQL database (short term memory)
     "pymysql==1.1.1", # For MySQL database (short term memory)
     "aiomysql==0.3.2", # For async MySQL database (short term memory)
     "opensearch-py==2.8.0",
@@ -51,6 +51,7 @@ extensions = [
     "redis>=5.0",                 # For Redis database
     "tos>=2.8.4",                   # For TOS storage and Viking DB
     "llama-index-vector-stores-redis>=0.6.1",
+    "cozeloop>=0.1.21", #For Cozeloop Prompt manager
 ]
 database = [
     "redis>=5.0",                 # For Redis database

{veadk_python-0.5.1 → veadk_python-0.5.2}/veadk/agent.py

@@ -85,6 +85,8 @@ class Agent(LlmAgent):
         short_term_memory (Optional[ShortTermMemory]): Session-based memory for temporary context.
         long_term_memory (Optional[LongTermMemory]): Cross-session memory for persistent user context.
         tracers (list[BaseTracer]): List of tracers used for telemetry and monitoring.
+        enable_authz (bool): Whether to enable agent authorization checks.
+        auto_save_session (bool): Whether to automatically save sessions to long-term memory.
     """
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
@@ -140,6 +142,8 @@ class Agent(LlmAgent):
 
     enable_authz: bool = False
 
+    auto_save_session: bool = False
+
     def model_post_init(self, __context: Any) -> None:
         super().model_post_init(None)  # for sub_agents init
 
@@ -258,6 +262,29 @@ class Agent(LlmAgent):
         if self.prompt_manager:
             self.instruction = self.prompt_manager.get_prompt
 
+        if self.auto_save_session:
+            if self.long_term_memory is None:
+                logger.warning(
+                    "auto_save_session is enabled, but long_term_memory is not initialized."
+                )
+            else:
+                from veadk.memory.save_session_callback import (
+                    save_session_to_long_term_memory,
+                )
+
+                if self.after_agent_callback:
+                    if isinstance(self.after_agent_callback, list):
+                        self.after_agent_callback.append(
+                            save_session_to_long_term_memory
+                        )
+                    else:
+                        self.after_agent_callback = [
+                            self.after_agent_callback,
+                            save_session_to_long_term_memory,
+                        ]
+                else:
+                    self.after_agent_callback = save_session_to_long_term_memory
+
         logger.info(f"VeADK version: {VERSION}")
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")

{veadk_python-0.5.1 → veadk_python-0.5.2}/veadk/cli/cli_rl.py

@@ -38,8 +38,8 @@ def rl_group():
     "--platform",
     "-p",
     required=True,
-    type=click.Choice(["ark"], case_sensitive=False),
-    help="Scaffold platform type (only support for now: ark)",
+    type=click.Choice(["ark", "lightning"], case_sensitive=False),
+    help="Scaffold platform type (supported: ark, lightning)",
 )
 @click.option(
     "--workspace", "-w", required=True, type=str, help="Target workspace directory name"
@@ -52,8 +52,9 @@ def rl_group():
 )
 def rl_init(platform: str, workspace: str, overwrite: bool):
     """
-    Initialize RL scaffold project for ark platform
+    Initialize RL scaffold project for ark or lightning platform
     Example: veadk rl init --platform ark --workspace veadk_rl_ark_project
+    Example: veadk rl init --platform lightning --workspace veadk_rl_lightning_project
     """
     # Locate template directory
     rl_template_root = get_rl_template_root()

veadk_python-0.5.2/veadk/cli/templates/rl/lightning/README.md

@@ -0,0 +1,58 @@
+# Agent Lightning
+
+Agent Lightning 提供了灵活且可扩展的框架，实现了智能体（client）和训练（server）的完全解耦。
+VeADK 与 Agent Lightning 集成，用户使用 VeADK 提供的脚手架，可以开发 VeADK Agent，然后运行 client 与 server 进行强化学习优化。
+
+## 准备工作
+
+在你的终端中运行以下命令，初始化一个 Agent Lightning 项目：
+
+```bash
+veadk rl init --platform lightning --workspace veadk_rl_lightning_project
+```
+
+该命令会在当前目录下创建一个名为 `veadk_rl_lightning_project` 的文件夹，其中包含了一个基本的基于 VeADK 和 Agent Lightning 的强化学习项目结构。
+然后在终端1中运行以下命令，启动 client：
+
+```bash
+cd veadk_rl_lightning_project
+python veadk_agent.py
+```
+
+然后在终端2中运行以下命令
+
+- 首先重启 ray 集群：
+
+```bash
+cd veadk_rl_lightning_project
+bash restart_ray.sh
+```  
+
+- 启动 server：
+
+```bash
+cd veadk_rl_lightning_project
+bash train.sh
+```
+
+## 原理说明
+
+生成后的项目结构如下，其中核心文件包括：
+
+- agent_client: `*_agent.py` 中定义了agent的rollout逻辑和reward规则
+- training_server: `train.sh` 定义了训练相关参数,用于启动训练服务器
+
+```shell
+veadk_rl_lightning_project
+├── data 
+    ├── demo_train.parquet # 训练数据,必须为 parquet 格式
+    ├── demo_test.parquet # 测试数据,必须为 parquet 格式
+└── demo_calculate_agent.py # agent的rollout逻辑和reward设定
+└── train.sh # 训练服务器启动脚本,设定训练相关参数 
+└── restart_ray.sh # 重启 ray 集群脚本
+```
+
+## 最佳实践案例
+
+1. 脚手架中，基于 VeADK 的算术 Agent 进行强化学习优化
+2. 启动 client (python demo_calculate_agent.py), 重启ray集群(bash restart_ray.sh), 最后启动训练服务器server (bash train.sh)，分别在终端1与终端2中运行以上命令

veadk_python-0.5.2/veadk/cli/templates/rl/lightning/demo_calculate_agent.py

@@ -0,0 +1,148 @@
+# 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 math
+import re
+import string
+import sympy
+from typing import Any, cast
+from veadk.agent import Agent
+from veadk.runner import Runner
+from veadk.memory.short_term_memory import ShortTermMemory
+from agentlightning import (
+    LLM,
+    LitAgent,
+    NamedResources,
+    Trainer,
+    reward,
+)
+
+
+def normalize_option(option: str) -> str:
+    """
+    >>> normalize_option("  (A)  \n")
+    'A'
+    """
+    return re.sub(r"(\s+|\(|\))", "", option)
+
+
+def is_option_result(result: str) -> bool:
+    """
+    >>> is_option_result("  A)  \n")
+    True
+    >>> is_option_result("  23/7 ")
+    False
+    """
+    return normalize_option(result) in list(string.ascii_letters)
+
+
+def float_eval(input_str: str) -> float:
+    if " = around " in input_str:
+        input_str = input_str.split(" = around ")[0]
+    expr = sympy.parse_expr(input_str, evaluate=True)
+    return float(expr.evalf())
+
+
+def scalar_are_results_same(pred_result: str, true_result: str, rel_tol: float) -> bool:
+    pred_result = str(pred_result) if pred_result is not None else ""  # type: ignore
+    true_result = str(true_result) if true_result is not None else ""  # type: ignore
+
+    if pred_result.strip() == true_result.strip():
+        return True
+
+    if is_option_result(true_result):
+        # The task is to select correct option
+        true_result = normalize_option(true_result)
+        pred_result = normalize_option(pred_result)
+        return pred_result == true_result
+
+    # The task is to calculate the result as a number
+    try:
+        pred_float = float_eval(pred_result)
+        true_float = float_eval(true_result)
+        return math.isclose(pred_float, true_float, rel_tol=rel_tol)
+    except Exception:
+        pass
+
+    return False
+
+
+@reward
+async def eval(prediction: str, ground_truth: str) -> float:
+    return float(scalar_are_results_same(prediction, ground_truth, 1e-2))
+
+
+class CalcAgent(LitAgent[Any]):
+    async def training_rollout_async(
+        self, task: Any, rollout_id: str, resources: NamedResources
+    ) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
+        calc_agent = Agent(
+            name="CalcAgent",
+            description="An agent that can perform calculations to answer questions.",
+            instruction="You are a helpful assistant that can perform mathematical calculations to answer questions accurately.",
+            model_provider="openai",
+            model=llm.model,
+            api_base=llm.endpoint,
+            api_key="",
+        )
+        runner = Runner(
+            agent=calc_agent,
+            short_term_memory=ShortTermMemory(),
+            app_name="calc_agent",
+            user_id="veadk_default_user",
+        )
+        try:
+            output_format = "Output the answer when you are ready. The answer should be surrounded by three sharps (`###`), in the form of ### ANSWER: <answer> ###."
+            prompt = task["question"] + " " + output_format
+            result = await runner.run(
+                session_id=rollout_id,
+                messages=prompt,
+            )
+            # evaluate
+            answer = re.search(
+                r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content
+            )  # type: ignore
+            if answer:
+                answer = answer.group(1)
+            else:
+                answer = result.messages[-1].content  # type: ignore
+        except Exception as e:
+            print("Failure:", str(e))
+            answer = "None"
+        reward = await eval(
+            answer, str(task["result"])
+        )  # reward is tracked with the decorator  # type: ignore
+        print(
+            "answer: {} ground_truth: {} reward: {}".format(
+                answer, task["result"], reward
+            )
+        )  # type: ignore
+
+    async def validation_rollout_async(
+        self, task: Any, rollout_id: str, resources: NamedResources
+    ) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
+        resources = {
+            "main_llm": LLM(
+                endpoint=llm.endpoint,
+                model=llm.model,
+                sampling_parameters={"temperature": 0},
+            )
+        }
+        return await self.training_rollout_async(task, rollout_id, resources)
+
+
+if __name__ == "__main__":
+    Trainer(n_workers=10).fit(CalcAgent(), "http://localhost:9999/")

veadk_python-0.5.2/veadk/cli/templates/rl/lightning/restart_ray.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+set -ex
+
+ray stop -v --force --grace-period 60
+ps aux
+env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head  --dashboard-host=0.0.0.0 --port 6380 --dashboard-port 8266

veadk_python-0.5.2/veadk/cli/templates/rl/lightning/train.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+set -e
+
+export N_GPUS=1
+export BASE_MODEL=Qwen/Qwen2.5-1.5B-Instruct
+export DATA_DIR=data
+export ROLLOUT_TP_SIZE=1
+export EXPERIMENT_NAME=calc_x
+export PROJECT_NAME=AgentLightning
+
+echo "Starting training script..."
+
+python -m agentlightning.verl \
+    algorithm.adv_estimator=grpo \
+    data.train_files=${DATA_DIR}/train.parquet \
+    data.val_files=${DATA_DIR}/test.parquet \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP_SIZE \
+    trainer.n_gpus_per_node=${N_GPUS} \
+    data.train_batch_size=32 \
+    actor_rollout_ref.rollout.n=4 \
+    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+    actor_rollout_ref.rollout.multi_turn.format=hermes \
+    actor_rollout_ref.model.path=${BASE_MODEL} \
+    data.max_prompt_length=4096 \
+    data.max_response_length=2048 \
+    data.truncation='error' \
+    trainer.val_before_train=True \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_remove_padding=True \
+    actor_rollout_ref.actor.use_kl_loss=False \
+    actor_rollout_ref.actor.kl_loss_coef=0.000 \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.actor.clip_ratio_low=0.2 \
+    actor_rollout_ref.actor.clip_ratio_high=0.3 \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.actor.fsdp_config.param_offload=True \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=True \
+    algorithm.use_kl_in_reward=False \
+    trainer.critic_warmup=0 \
+    trainer.logger=['console','wandb'] \
+    trainer.project_name=${PROJECT_NAME} \
+    trainer.experiment_name=${EXPERIMENT_NAME} \
+    trainer.nnodes=1 \
+    trainer.save_freq=256 \
+    trainer.test_freq=32 \
+    trainer.total_epochs=2 $@

{veadk_python-0.5.1 → veadk_python-0.5.2}/veadk/consts.py

@@ -65,7 +65,7 @@ DEFAULT_COZELOOP_SPACE_NAME = "VeADK Space"
 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_NAME = "doubao-seedance-1-5-pro-251215"
 DEFAULT_VIDEO_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
 
 DEFAULT_IMAGE_GENERATE_MODEL_NAME = "doubao-seedream-4-5-251128"

veadk_python-0.5.2/veadk/memory/save_session_callback.py

@@ -0,0 +1,159 @@
+# 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 time
+from google.adk.agents.callback_context import CallbackContext
+from veadk.config import getenv
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+# Session-level cache for tracking save state
+# Format: {(app_name, user_id, session_id): {'last_save_time': float, 'last_event_count': int}}
+_session_save_cache: dict = {}
+
+# Track active session per user to detect session switches
+# Format: {(app_name, user_id): session_id}
+_active_sessions: dict = {}
+
+# Configurable thresholds
+MIN_MESSAGES_THRESHOLD = getenv(
+    "MIN_MESSAGES_THRESHOLD", 10
+)  # Minimum number of new messages before saving
+MIN_TIME_THRESHOLD = getenv(
+    "MIN_TIME_THRESHOLD", 60
+)  # Minimum seconds between saves (1 minute)
+
+
+async def save_session_to_long_term_memory(
+    callback_context: CallbackContext,
+) -> None:
+    """Save the current session to long-term memory.
+
+    Args:
+        callback_context: The callback context containing invocation information.
+
+    Returns:
+        None
+    """
+    try:
+        agent = callback_context._invocation_context.agent
+
+        long_term_memory = getattr(agent, "long_term_memory", None)
+        if not long_term_memory:
+            logger.error(
+                "Long-term memory is not initialized in agent, cannot save session to memory."
+            )
+            return None
+
+        app_name = callback_context._invocation_context.app_name
+        user_id = callback_context._invocation_context.user_id
+        session_id = callback_context._invocation_context.session.id
+        session_service = callback_context._invocation_context.session_service
+
+        current_time = time.time()
+
+        # Detect session switch and force save previous session
+        user_key = (app_name, user_id)
+        previous_session_id = _active_sessions.get(user_key)
+
+        if previous_session_id and previous_session_id != session_id:
+            logger.info(
+                f"Session switch detected for user {user_id}: "
+                f"{previous_session_id} -> {session_id}. "
+                f"Force saving previous session."
+            )
+            old_session = await session_service.get_session(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=previous_session_id,
+            )
+            if old_session:
+                old_events = getattr(old_session, "events", [])
+                old_event_count = len(old_events)
+                await long_term_memory.add_session_to_memory(old_session)
+                old_cache_key = (app_name, user_id, previous_session_id)
+
+                _session_save_cache[old_cache_key] = {
+                    "last_save_time": current_time,
+                    "last_event_count": old_event_count,
+                }
+                logger.info(
+                    f"Previous session `{old_session.id}` saved to long term memory due to session switch."
+                )
+
+        # Update active session
+        _active_sessions[user_key] = session_id
+
+        session = await session_service.get_session(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+        )
+
+        if not session:
+            logger.error(
+                f"Session {session_id} (app_name={app_name}, user_id={user_id}) not found in session service, cannot save to long-term memory."
+            )
+            return None
+
+        current_events = getattr(session, "events", [])
+        current_event_count = len(current_events)
+        # logger.debug(f"Current event count: {current_event_count}")
+
+        # Create cache key
+        cache_key = (app_name, user_id, session_id)
+
+        cache_info = _session_save_cache.get(cache_key)
+
+        if cache_info:
+            last_save_time = cache_info.get("last_save_time", 0)
+            last_event_count = cache_info.get("last_event_count", 0)
+
+            time_elapsed = current_time - last_save_time
+            new_events_count = current_event_count - last_event_count
+
+            # Check if we should skip save
+            if (
+                time_elapsed < MIN_TIME_THRESHOLD
+                and new_events_count < MIN_MESSAGES_THRESHOLD
+            ):
+                logger.info(
+                    f"Skipping save for session {session_id}: "
+                    f"only {new_events_count} new events (need {MIN_MESSAGES_THRESHOLD}) "
+                    f"and {time_elapsed:.1f}s elapsed (need {MIN_TIME_THRESHOLD}s)"
+                )
+                return None
+        else:
+            logger.info(f"First save for session {session_id}.")
+
+        # Save to long-term memory
+        await long_term_memory.add_session_to_memory(session)
+
+        # Update cache
+        _session_save_cache[cache_key] = {
+            "last_save_time": current_time,
+            "last_event_count": current_event_count,
+        }
+
+        logger.info(f"Add session `{session.id}` to long term memory.")
+
+        return None
+
+    except AttributeError as e:
+        logger.error(f"AttributeError while saving session to memory: {e}")
+        return None
+    except Exception as e:
+        logger.error(f"Unexpected error while saving session to memory: {e}")
+        return None

{veadk_python-0.5.1 → veadk_python-0.5.2}/veadk/tools/builtin_tools/image_generate.py

@@ -83,7 +83,10 @@ def _build_input_parts(item: dict, task_type: str, image_field):
 
 
 def handle_single_task_sync(
-    idx: int, item: dict, tool_context
+    idx: int,
+    item: dict,
+    timeout: int,
+    tool_context,
 ) -> tuple[list[dict], list[str]]:
     logger.debug(f"handle_single_task_sync item {idx}: {item}")
     success_list: list[dict] = []
@@ -139,6 +142,7 @@ def handle_single_task_sync(
                             "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
                         ),
                     },
+                    timeout=timeout,
                 )
             else:
                 response = client.images.generate(
@@ -152,6 +156,7 @@ def handle_single_task_sync(
                             "MODEL_AGENT_CLIENT_REQ_ID", f"veadk/{VERSION}"
                         ),
                     },
+                    timeout=timeout,
                 )
 
             if not response.error:
@@ -228,14 +233,16 @@ def handle_single_task_sync(
     return success_list, error_list
 
 
-async def image_generate(tasks: list[dict], tool_context) -> Dict:
-    """Generate images with Seedream 4.0.
+async def image_generate(tasks: list[dict], tool_context, timeout: int = 600) -> Dict:
+    """Generate images with Seedream 4.0 / 4.5
 
     Commit batch image generation requests via tasks.
 
     Args:
         tasks (list[dict]):
             A list of image-generation tasks. Each task is a dict.
+        timeout (int)
+            The timeout limit for the image generation task request, in seconds, with a default value of 600 seconds.
     Per-task schema
     ---------------
     Required:
@@ -336,7 +343,9 @@ async def image_generate(tasks: list[dict], tool_context) -> Dict:
 
         def make_task(idx, item):
             ctx = base_ctx.copy()
-            return lambda: ctx.run(handle_single_task_sync, idx, item, tool_context)
+            return lambda: ctx.run(
+                handle_single_task_sync, idx, item, timeout, tool_context
+            )
 
         loop = asyncio.get_event_loop()
         futures = [