fix: transport response_id by session state

Author: zakahan

veadk/agent.py

@@ -36,7 +36,7 @@
 from veadk.knowledgebase import KnowledgeBase
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
-from veadk.models.ark_llm import add_previous_response_id
+from veadk.models.ark_llm import add_previous_response_id, add_response_id
 from veadk.processors import BaseRunProcessor, NoOpRunProcessor
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
@@ -213,6 +213,10 @@ def model_post_init(self, __context: Any) -> None:
                     self.before_model_callback = add_previous_response_id
                 else:
                     self.before_model_callback.append(add_previous_response_id)
+                if not self.after_tool_callback:
+                    self.after_model_callback = add_response_id
+                else:
+                    self.after_model_callback.append(add_response_id)
             else:
                 self.model = LiteLlm(
                     model=f"{self.model_provider}/{self.model_name}",

veadk/models/ark_llm.py

@@ -259,13 +259,18 @@ def add_previous_response_id(
 ) -> Optional[LlmResponse]:
     invocation_context = callback_context._invocation_context
     events = invocation_context.session.events
+    agent_name = callback_context.agent_name
+    # read response_id
     if (
         events
         and len(events) >= 2
         and events[-2].custom_metadata
         and "response_id" in events[-2].custom_metadata
     ):
-        previous_response_id = events[-2].custom_metadata["response_id"]
+        previous_response_id = callback_context.state.get(
+            f"agent:{agent_name}:response_id"
+        )
+        # previous_response_id = events[-2].custom_metadata["response_id"]
         if "contents_count" in CacheMetadata.model_fields:  # adk >= 1.17
             llm_request.cache_metadata = CacheMetadata(
                 cache_name=previous_response_id,
@@ -283,3 +288,14 @@ def add_previous_response_id(
                 cached_contents_count=0,
             )
     return
+
+
+# after_model_callback
+def add_response_id(
+    callback_context: CallbackContext, llm_response: LlmResponse
+) -> Optional[LlmResponse]:
+    agent_name = callback_context.agent_name
+
+    response_id = llm_response.custom_metadata["response_id"]
+    callback_context.state[f"agent:{agent_name}:response_id"] = response_id
+    return