fix: remove before_model_callback

Author: zakahan

veadk/agent.py

@@ -19,9 +19,10 @@
 
 from google.adk.agents import LlmAgent, RunConfig, InvocationContext
 from google.adk.agents.base_agent import BaseAgent
+from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.agents.llm_agent import InstructionProvider, ToolUnion
 from google.adk.agents.run_config import StreamingMode
-from google.adk.events import Event, EventActions
+from google.adk.events import Event
 from google.adk.models.lite_llm import LiteLlm
 from google.adk.runners import Runner
 from google.genai import types
@@ -37,7 +38,6 @@
 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.processors import BaseRunProcessor, NoOpRunProcessor
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
@@ -155,6 +155,8 @@ class Agent(LlmAgent):
 
     enable_responses: bool = False
 
+    context_cache_config: Optional[ContextCacheConfig] = None
+
     run_processor: Optional[BaseRunProcessor] = Field(default=None, exclude=True)
     """Optional run processor for intercepting and processing agent execution flows.
 
@@ -210,16 +212,12 @@ def model_post_init(self, __context: Any) -> None:
                     api_base=self.model_api_base,
                     **self.model_extra_config,
                 )
-                if not self.before_model_callback:
-                    self.before_model_callback = add_previous_response_id
-                else:
-                    if isinstance(self.before_model_callback, list):
-                        self.before_model_callback.append(add_previous_response_id)
-                    else:
-                        self.before_model_callback = [
-                            self.before_model_callback,
-                            add_previous_response_id,
-                        ]
+                if not self.context_cache_config:
+                    self.context_cache_config = ContextCacheConfig(
+                        cache_intervals=100,  # maximum number
+                        ttl_seconds=315360000,
+                        min_tokens=0,
+                    )
             else:
                 self.model = LiteLlm(
                     model=f"{self.model_provider}/{self.model_name}",
@@ -265,22 +263,11 @@ def model_post_init(self, __context: Any) -> None:
     async def _run_async_impl(
         self, ctx: InvocationContext
     ) -> AsyncGenerator[Event, None]:
+        if self.enable_responses:
+            if not ctx.context_cache_config:
+                ctx.context_cache_config = self.context_cache_config
+
         async for event in super()._run_async_impl(ctx):
-            agent_name = self.name
-            if (
-                self.enable_responses
-                and event.custom_metadata
-                and event.custom_metadata.get("response_id")
-            ):
-                response_id = event.custom_metadata["response_id"]
-                yield Event(
-                    invocation_id=ctx.invocation_id,
-                    author=self.name,
-                    actions=EventActions(
-                        state_delta={f"agent:{agent_name}:response_id": response_id}
-                    ),
-                    branch=ctx.branch,
-                )
             yield event
 
     async def _run(

veadk/models/ark_llm.py

@@ -15,13 +15,11 @@
 # adapted from Google ADK models adk-python/blob/main/src/google/adk/models/lite_llm.py at f1f44675e4a86b75e72cfd838efd8a0399f23e24 · google/adk-python
 
 import json
-from typing import Any, Dict, Union, AsyncGenerator, Optional
+from typing import Any, Dict, Union, AsyncGenerator
 
 import litellm
 import openai
 from openai.types.responses import Response as OpenAITypeResponse, ResponseStreamEvent
-from google.adk.agents.callback_context import CallbackContext
-from google.adk.models.cache_metadata import CacheMetadata
 from google.adk.models import LlmRequest, LlmResponse
 from google.adk.models.lite_llm import (
     LiteLlm,
@@ -41,6 +39,7 @@
 
 from veadk.models.ark_transform import (
     CompletionToResponsesAPIHandler,
+    get_previous_response_id,
 )
 from veadk.utils.logger import get_logger
 
@@ -90,7 +89,7 @@ async def generate_content_async(
         Yields:
           LlmResponse: The model response.
         """
-
+        agent_name = llm_request.config.labels["adk_agent_name"]
         self._maybe_append_user_content(llm_request)
         # logger.debug(_build_request_log(llm_request))
 
@@ -105,7 +104,10 @@ async def generate_content_async(
         # get previous_response_id
         previous_response_id = None
         if llm_request.cache_metadata and llm_request.cache_metadata.cache_name:
-            previous_response_id = llm_request.cache_metadata.cache_name
+            previous_response_id = get_previous_response_id(
+                llm_request.cache_metadata,
+                agent_name,
+            )
         completion_args = {
             "model": self.model,
             "messages": messages,
@@ -210,6 +212,7 @@ async def generate_content_async(
                             )
                         )
                         self.transform_handler.adapt_responses_api(
+                            llm_request,
                             model_response,
                             aggregated_llm_response_with_tool_call,
                             stream=True,
@@ -223,7 +226,10 @@ async def generate_content_async(
                             )
                         )
                         self.transform_handler.adapt_responses_api(
-                            model_response, aggregated_llm_response, stream=True
+                            llm_request,
+                            model_response,
+                            aggregated_llm_response,
+                            stream=True,
                         )
                         text = ""
 
@@ -248,32 +254,6 @@ async def generate_content_async(
             for (
                 llm_response
             ) in self.transform_handler.openai_response_to_generate_content_response(
-                raw_response
+                llm_request, raw_response
             ):
                 yield llm_response
-
-
-# before_model_callback
-def add_previous_response_id(
-    callback_context: CallbackContext, llm_request: LlmRequest
-) -> Optional[LlmResponse]:
-    agent_name = callback_context.agent_name
-    # read response_id
-    previous_response_id = callback_context.state.get(f"agent:{agent_name}:response_id")
-    if "contents_count" in CacheMetadata.model_fields:  # adk >= 1.17
-        llm_request.cache_metadata = CacheMetadata(
-            cache_name=previous_response_id,
-            expire_time=0,
-            fingerprint="",
-            invocations_used=0,
-            contents_count=0,
-        )
-    else:  # 1.15 <= adk < 1.17
-        llm_request.cache_metadata = CacheMetadata(
-            cache_name=previous_response_id,
-            expire_time=0,
-            fingerprint="",
-            invocations_used=0,
-            cached_contents_count=0,
-        )
-    return

veadk/models/ark_transform.py

@@ -14,11 +14,13 @@
 
 # adapted from Google ADK models adk-python/blob/main/src/google/adk/models/lite_llm.py at f1f44675e4a86b75e72cfd838efd8a0399f23e24 · google/adk-python
 
+import json
 import uuid
 from typing import Any, Dict, Optional, cast, List, Generator, Tuple, Union
 
 import litellm
-from google.adk.models import LlmResponse
+from google.adk.models import LlmResponse, LlmRequest
+from google.adk.models.cache_metadata import CacheMetadata
 from google.adk.models.lite_llm import (
     TextChunk,
     FunctionChunk,
@@ -139,6 +141,74 @@ def ark_field_reorganization(request_data: dict) -> dict:
     return request_data
 
 
+def build_cache_metadata(agent_response_id: dict) -> CacheMetadata:
+    """Create a new CacheMetadata instance for agent response tracking.
+
+    Args:
+        agent_name: Name of the agent
+        response_id: Response ID to track
+
+    Returns:
+        A new CacheMetadata instance with the agent-response mapping
+    """
+    cache_name = json.dumps(agent_response_id)
+    if "contents_count" in CacheMetadata.model_fields:  # adk >= 1.17
+        cache_metadata = CacheMetadata(
+            cache_name=cache_name,
+            expire_time=0,
+            fingerprint="",
+            invocations_used=0,
+            contents_count=0,
+        )
+    else:  # 1.15 <= adk < 1.17
+        cache_metadata = CacheMetadata(
+            cache_name=cache_name,
+            expire_time=0,
+            fingerprint="",
+            invocations_used=0,
+            cached_contents_count=0,
+        )
+    return cache_metadata
+
+
+def update_cache_metadata(
+    cache_metadata: CacheMetadata,
+    agent_name: str,
+    response_id: str,
+) -> CacheMetadata:
+    """Update cache metadata by creating a new instance with updated cache_name.
+
+    Since CacheMetadata is frozen, we cannot modify it directly. Instead,
+    we create a new instance with the updated cache_name field.
+    """
+    try:
+        agent_response_id = json.loads(cache_metadata.cache_name)
+        agent_response_id[agent_name] = response_id
+        updated_cache_name = agent_response_id
+
+        # Create a new CacheMetadata instance with updated cache_name
+        return build_cache_metadata(updated_cache_name)
+    except json.JSONDecodeError as e:
+        logger.warning(
+            f"Failed to update cache metadata. The cache_name is not a valid JSON string., {str(e)}"
+        )
+        return cache_metadata
+
+
+def get_previous_response_id(
+    cache_metadata: CacheMetadata,
+    agent_name: str,
+):
+    try:
+        agent_response_id = json.loads(cache_metadata.cache_name)
+        return agent_response_id.get(agent_name, None)
+    except json.JSONDecodeError as e:
+        logger.warning(
+            f"Failed to get previous response id. The cache_name is not a valid JSON string., {str(e)}"
+        )
+        return None
+
+
 class CompletionToResponsesAPIHandler:
     def __init__(self):
         self.litellm_handler = LiteLLMResponsesTransformationHandler()
@@ -231,7 +301,7 @@ def transform_response(
         return result_list
 
     def openai_response_to_generate_content_response(
-        self, raw_response: OpenAITypeResponse
+        self, llm_request: LlmRequest, raw_response: OpenAITypeResponse
     ) -> list[LlmResponse]:
         """
         OpenAITypeResponse -> litellm.ModelResponse -> LlmResponse
@@ -246,6 +316,7 @@ def openai_response_to_generate_content_response(
             llm_response = _model_response_to_generate_content_response(model_response)
 
             llm_response = self.adapt_responses_api(
+                llm_request,
                 model_response,
                 llm_response,
             )
@@ -254,6 +325,7 @@ def openai_response_to_generate_content_response(
 
     def adapt_responses_api(
         self,
+        llm_request: LlmRequest,
         model_response: ModelResponse,
         llm_response: LlmResponse,
         stream: bool = False,
@@ -262,9 +334,21 @@ def adapt_responses_api(
         Adapt responses api.
         """
         if not model_response.id.startswith("chatcmpl"):
-            if llm_response.custom_metadata is None:
-                llm_response.custom_metadata = {}
-            llm_response.custom_metadata["response_id"] = model_response["id"]
+            # if llm_response.custom_metadata is None:
+            #     llm_response.custom_metadata = {}
+            # llm_response.custom_metadata["response_id"] = model_response["id"]
+            previous_response_id = model_response["id"]
+            if not llm_request.cache_metadata:
+                llm_response.cache_metadata = build_cache_metadata(
+                    {llm_request.config.labels["adk_agent_name"]: previous_response_id}
+                )
+            else:
+                llm_response.cache_metadata = update_cache_metadata(
+                    llm_request.cache_metadata,
+                    llm_request.config.labels["adk_agent_name"],
+                    previous_response_id,
+                )
+
         # add responses cache data
         if not stream:
             if model_response.get("usage", {}).get("prompt_tokens_details"):