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

{veadk_python-0.2.35 → veadk_python-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.2.35
+Version: 0.5.1
 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
@@ -236,7 +236,7 @@ Requires-Dist: aiomysql==0.3.2
 Requires-Dist: opensearch-py==2.8.0
 Requires-Dist: filetype==1.2.0
 Requires-Dist: vikingdb-python-sdk==0.1.3
-Requires-Dist: agentkit-sdk-python==0.2.0
+Requires-Dist: agentkit-sdk-python>=0.2.0
 Provides-Extra: extensions
 Requires-Dist: redis>=5.0; extra == "extensions"
 Requires-Dist: tos>=2.8.4; extra == "extensions"

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

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.35"
+version = "0.5.1"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -40,7 +40,7 @@ dependencies = [
     "opensearch-py==2.8.0",
     "filetype==1.2.0",
     "vikingdb-python-sdk==0.1.3",
-    "agentkit-sdk-python==0.2.0"
+    "agentkit-sdk-python>=0.2.0"
 ]
 
 [project.scripts]
@@ -73,7 +73,6 @@ dev = [
     "pytest-xdist>=3.8.0",
 ]
 
-
 [dependency-groups]
 dev = [
     "pre-commit>=4.2.0",            # Format checking

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

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import os
-from typing import Optional, Union
+from typing import Optional, Union, AsyncGenerator
 
 # If user didn't set LITELLM_LOCAL_MODEL_COST_MAP, set it to True
 # to enable local model cost map.
@@ -24,10 +24,12 @@ from typing import Optional, Union
 if not os.getenv("LITELLM_LOCAL_MODEL_COST_MAP"):
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
 
-from google.adk.agents import LlmAgent, RunConfig
+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.models.lite_llm import LiteLlm
 from google.adk.runners import Runner
 from google.genai import types
@@ -44,10 +46,15 @@ from veadk.knowledgebase import KnowledgeBase
 from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.processors import BaseRunProcessor, NoOpRunProcessor
-from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
+from veadk.prompts.agent_default_prompt import (
+    DEFAULT_DESCRIPTION,
+    DEFAULT_INSTRUCTION,
+)
+from veadk.prompts.prompt_manager import BasePromptManager
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_asyncio, patch_tracer
+from veadk.utils.misc import check_litellm_version
 from veadk.version import VERSION
 
 patch_tracer()
@@ -67,7 +74,7 @@ class Agent(LlmAgent):
         name (str): The name of the agent.
         description (str): A description of the agent, useful in A2A scenarios.
         instruction (Union[str, InstructionProvider]): The instruction or instruction provider.
-        model_name (str): Name of the model used by the agent.
+        model_name (Union[str, List[str]]): Name of the model used by the agent.
         model_provider (str): Provider of the model (e.g., openai).
         model_api_base (str): The base URL of the model API.
         model_api_key (str): The API key for accessing the model.
@@ -86,7 +93,9 @@ class Agent(LlmAgent):
     description: str = DEFAULT_DESCRIPTION
     instruction: Union[str, InstructionProvider] = DEFAULT_INSTRUCTION
 
-    model_name: str = Field(default_factory=lambda: settings.model.name)
+    model_name: Union[str, list[str]] = Field(
+        default_factory=lambda: settings.model.name
+    )
     model_provider: str = Field(default_factory=lambda: settings.model.provider)
     model_api_base: str = Field(default_factory=lambda: settings.model.api_base)
     model_api_key: str = Field(default_factory=lambda: settings.model.api_key)
@@ -96,6 +105,8 @@ class Agent(LlmAgent):
 
     sub_agents: list[BaseAgent] = Field(default_factory=list, exclude=True)
 
+    prompt_manager: Optional[BasePromptManager] = None
+
     knowledgebase: Optional[KnowledgeBase] = None
 
     short_term_memory: Optional[ShortTermMemory] = None
@@ -103,6 +114,10 @@ class Agent(LlmAgent):
 
     tracers: list[BaseTracer] = []
 
+    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.
 
@@ -151,12 +166,50 @@ class Agent(LlmAgent):
         logger.info(f"Model extra config: {self.model_extra_config}")
 
         if not self.model:
-            self.model = LiteLlm(
-                model=f"{self.model_provider}/{self.model_name}",
-                api_key=self.model_api_key,
-                api_base=self.model_api_base,
-                **self.model_extra_config,
-            )
+            if self.enable_responses:
+                min_version = "1.79.3"
+                check_litellm_version(min_version)
+
+                from veadk.models.ark_llm import ArkLlm
+
+                self.model = ArkLlm(
+                    model=f"{self.model_provider}/{self.model_name}",
+                    api_key=self.model_api_key,
+                    api_base=self.model_api_base,
+                    **self.model_extra_config,
+                )
+                if not self.context_cache_config:
+                    self.context_cache_config = ContextCacheConfig(
+                        cache_intervals=100,  # maximum number
+                        ttl_seconds=315360000,
+                        min_tokens=0,
+                    )
+            else:
+                fallbacks = None
+                if isinstance(self.model_name, list):
+                    if self.model_name:
+                        model_name = self.model_name[0]
+                        fallbacks = [
+                            f"{self.model_provider}/{m}" for m in self.model_name[1:]
+                        ]
+                        logger.info(
+                            f"Using primary model: {model_name}, with fallbacks: {self.model_name[1:]}"
+                        )
+                    else:
+                        model_name = settings.model.name
+                        logger.warning(
+                            f"Empty model_name list provided, using default model from settings: {model_name}"
+                        )
+                else:
+                    model_name = self.model_name
+
+                self.model = LiteLlm(
+                    model=f"{self.model_provider}/{model_name}",
+                    api_key=self.model_api_key,
+                    api_base=self.model_api_base,
+                    fallbacks=fallbacks,
+                    **self.model_extra_config,
+                )
             logger.debug(
                 f"LiteLLM client created with config: {self.model_extra_config}"
             )
@@ -202,6 +255,9 @@ class Agent(LlmAgent):
             else:
                 self.before_agent_callback = check_agent_authorization
 
+        if self.prompt_manager:
+            self.instruction = self.prompt_manager.get_prompt
+
         logger.info(f"VeADK version: {VERSION}")
 
         logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
@@ -209,6 +265,28 @@ class Agent(LlmAgent):
             f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools'})}"
         )
 
+    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):
+            yield event
+            if self.enable_responses and event.cache_metadata:
+                # for persistent short-term memory with response api
+                session_state_event = Event(
+                    invocation_id=event.invocation_id,
+                    author=event.author,
+                    actions=EventActions(
+                        state_delta={
+                            "response_id": event.cache_metadata.cache_name,
+                        }
+                    ),
+                )
+                yield session_state_event
+
     async def _run(
         self,
         runner,
@@ -274,14 +352,20 @@ class Agent(LlmAgent):
             return
 
         if not self.tracers:
-            from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
+            from veadk.tracing.telemetry.opentelemetry_tracer import (
+                OpentelemetryTracer,
+            )
 
             self.tracers.append(OpentelemetryTracer())
 
         exporters = self.tracers[0].exporters  # type: ignore
 
-        from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
-        from veadk.tracing.telemetry.exporters.cozeloop_exporter import CozeloopExporter
+        from veadk.tracing.telemetry.exporters.apmplus_exporter import (
+            APMPlusExporter,
+        )
+        from veadk.tracing.telemetry.exporters.cozeloop_exporter import (
+            CozeloopExporter,
+        )
         from veadk.tracing.telemetry.exporters.tls_exporter import TLSExporter
 
         if enable_apmplus_tracer and not any(

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/configs/database_configs.py

@@ -28,6 +28,10 @@ class OpensearchConfig(BaseSettings):
 
     port: int = 9200
 
+    cert_path: str = ""
+
+    use_ssl: bool = True
+
     username: str = ""
 
     password: str = ""

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/requirements.txt

@@ -1,4 +1,4 @@
 Flask==2.3.2
 Flask-SQLAlchemy==3.0.5
-Werkzeug==2.3.6
-gunicorn==20.1.0
+werkzeug>=3.1.4
+gunicorn>=22.0.0

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/integrations/ve_viking_db_memory/ve_viking_db_memory.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 
 import json
+import os
 import threading
-from veadk.utils.misc import getenv
+
 from volcengine.ApiInfo import ApiInfo
 from volcengine.auth.SignerV4 import SignerV4
 from volcengine.base.Service import Service
 from volcengine.Credentials import Credentials
 from volcengine.ServiceInfo import ServiceInfo
 
+from veadk.utils.misc import getenv
+
 
 class VikingDBMemoryException(Exception):
     def __init__(self, code, request_id, message=None):
@@ -56,7 +59,9 @@ class VikingDBMemoryClient(Service):
         socket_timeout=30,
     ):
         env_host = getenv(
-            "DATABASE_VIKINGMEM_BASE_URL", default_value=None, allow_false_values=True
+            "DATABASE_VIKINGMEM_BASE_URL",
+            default_value=None,
+            allow_false_values=True,
         )
         if env_host:
             if env_host.startswith("http://"):
@@ -85,7 +90,9 @@ class VikingDBMemoryClient(Service):
             self.get_body("Ping", {}, json.dumps({}))
         except Exception as e:
             raise VikingDBMemoryException(
-                1000028, "missed", "host or region is incorrect: {}".format(str(e))
+                1000028,
+                "missed",
+                "host or region is incorrect: {}".format(str(e)),
             ) from None
 
     def setHeader(self, header):
@@ -118,49 +125,70 @@ class VikingDBMemoryClient(Service):
                 "/api/memory/collection/create",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "GetCollection": ApiInfo(
                 "POST",
                 "/api/memory/collection/info",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "DropCollection": ApiInfo(
                 "POST",
                 "/api/memory/collection/delete",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "UpdateCollection": ApiInfo(
                 "POST",
                 "/api/memory/collection/update",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "SearchMemory": ApiInfo(
                 "POST",
                 "/api/memory/search",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "AddMessages": ApiInfo(
                 "POST",
                 "/api/memory/messages/add",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
             "Ping": ApiInfo(
                 "GET",
                 "/api/memory/ping",
                 {},
                 {},
-                {"Accept": "application/json", "Content-Type": "application/json"},
+                {
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
             ),
         }
         return api_info
@@ -199,7 +227,9 @@ class VikingDBMemoryClient(Service):
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception as e:
                 raise VikingDBMemoryException(
-                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                    1000028,
+                    "missed",
+                    "json load res error, res:{}".format(str(e)),
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
@@ -223,7 +253,9 @@ class VikingDBMemoryClient(Service):
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception as e:
                 raise VikingDBMemoryException(
-                    1000028, "missed", "json load res error, res:{}".format(str(e))
+                    1000028,
+                    "missed",
+                    "json load res error, res:{}".format(str(e)),
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
@@ -241,6 +273,7 @@ class VikingDBMemoryClient(Service):
         self,
         collection_name,
         description="",
+        project="default",
         custom_event_type_schemas=[],
         custom_entity_type_schemas=[],
         builtin_event_types=[],
@@ -248,6 +281,10 @@ class VikingDBMemoryClient(Service):
     ):
         params = {
             "CollectionName": collection_name,
+            "ProjectName": project,
+            "CollectionType": os.getenv(
+                "DATABASE_VIKINGMEM_COLLECTION_TYPE", "standard"
+            ),
             "Description": description,
             "CustomEventTypeSchemas": custom_event_type_schemas,
             "CustomEntityTypeSchemas": custom_entity_type_schemas,
@@ -257,8 +294,8 @@ class VikingDBMemoryClient(Service):
         res = self.json("CreateCollection", {}, json.dumps(params))
         return json.loads(res)
 
-    def get_collection(self, collection_name):
-        params = {"CollectionName": collection_name}
+    def get_collection(self, collection_name, project="default"):
+        params = {"CollectionName": collection_name, "ProjectName": project}
         res = self.json("GetCollection", {}, json.dumps(params))
         return json.loads(res)
 

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/knowledgebase/backends/opensearch_backend.py

@@ -27,9 +27,13 @@ from typing_extensions import Any, override
 
 import veadk.config  # noqa E401
 from veadk.configs.database_configs import OpensearchConfig
-from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.configs.model_configs import (
+    EmbeddingModelConfig,
+    NormalEmbeddingModelConfig,
+)
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.utils.logger import get_logger
 
 try:
     from llama_index.vector_stores.opensearch import (
@@ -42,6 +46,9 @@ except ImportError:
     )
 
 
+logger = get_logger(__name__)
+
+
 class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
     """Opensearch-based backend for knowledgebase.
 
@@ -79,6 +86,12 @@ class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
 
     def model_post_init(self, __context: Any) -> None:
         self.precheck_index_naming()
+
+        if not self.opensearch_config.cert_path:
+            logger.warning(
+                "OpenSearch cert_path is not set, which may lead to security risks"
+            )
+
         self._opensearch_client = OpensearchVectorClient(
             endpoint=self.opensearch_config.host,
             port=self.opensearch_config.port,
@@ -86,8 +99,9 @@ class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
                 self.opensearch_config.username,
                 self.opensearch_config.password,
             ),
-            use_ssl=True,
-            verify_certs=False,
+            use_ssl=self.opensearch_config.use_ssl,
+            verify_certs=False if not self.opensearch_config.cert_path else True,
+            ca_certs=self.opensearch_config.cert_path,
             dim=self.embedding_config.dim,
             index=self.index,  # collection name
         )

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -117,6 +117,10 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
         default_factory=lambda: os.getenv("DATABASE_VIKING_PROJECT", "default")
     )
 
+    version: str = Field(
+        default_factory=lambda: os.getenv("DATABASE_VIKING_VERSION", "2")
+    )
+
     region: str = Field(
         default_factory=lambda: os.getenv("DATABASE_VIKING_REGION", "cn-beijing")
     )
@@ -457,12 +461,16 @@ class VikingDBKnowledgeBackend(BaseKnowledgebaseBackend):
 
     def create_collection(self) -> None:
         CREATE_COLLECTION_PATH = "/api/knowledge/collection/create"
-
+        if self.version not in ["2", "4"]:
+            raise ValueError(
+                f"The version number must be 2 or 4. The current value: {self.version}. For details, please refer to: `https://www.volcengine.com/docs/84313/1254593?lang=zh`"
+            )
         response = self._do_request(
             body={
                 "name": self.index,
                 "project": self.volcengine_project,
                 "description": "Created by Volcengine Agent Development Kit (VeADK).",
+                "version": int(self.version),
             },
             path=CREATE_COLLECTION_PATH,
             method="POST",

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -22,7 +22,10 @@ from typing_extensions import Any, override
 
 import veadk.config  # noqa E401
 from veadk.configs.database_configs import OpensearchConfig
-from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
+from veadk.configs.model_configs import (
+    EmbeddingModelConfig,
+    NormalEmbeddingModelConfig,
+)
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
@@ -74,6 +77,11 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
 
         self.precheck_index_naming(index)
 
+        if not self.opensearch_config.cert_path:
+            logger.warning(
+                "OpenSearch cert_path is not set, which may lead to security risks"
+            )
+
         opensearch_client = OpensearchVectorClient(
             endpoint=self.opensearch_config.host,
             port=self.opensearch_config.port,
@@ -81,8 +89,9 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
                 self.opensearch_config.username,
                 self.opensearch_config.password,
             ),
-            use_ssl=True,
-            verify_certs=False,
+            use_ssl=self.opensearch_config.use_ssl,
+            verify_certs=False if not self.opensearch_config.cert_path else True,
+            ca_certs=self.opensearch_config.cert_path,
             dim=self.embedding_config.dim,
             index=index,
         )

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -21,6 +21,8 @@ from typing import Any
 
 from pydantic import Field
 from typing_extensions import override
+from vikingdb import IAM
+from vikingdb.memory import VikingMem
 
 import veadk.config  # noqa E401
 from veadk.auth.veauth.utils import get_credential_from_vefaas_iam
@@ -30,9 +32,6 @@ from veadk.integrations.ve_viking_db_memory.ve_viking_db_memory import (
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
-from vikingdb import IAM
-from vikingdb.memory import VikingMem
-
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -49,9 +48,16 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
 
     session_token: str = ""
 
-    region: str = "cn-beijing"
+    region: str = Field(
+        default_factory=lambda: os.getenv("DATABASE_VIKINGMEM_REGION") or "cn-beijing"
+    )
     """VikingDB memory region"""
 
+    volcengine_project: str = Field(
+        default_factory=lambda: os.getenv("DATABASE_VIKINGMEM_PROJECT") or "default"
+    )
+    """VikingDB memory project"""
+
     memory_type: list[str] = Field(default_factory=list)
 
     def model_post_init(self, __context: Any) -> None:
@@ -87,7 +93,9 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
     def _collection_exist(self) -> bool:
         try:
             client = self._get_client()
-            client.get_collection(collection_name=self.index)
+            client.get_collection(
+                collection_name=self.index, project=self.volcengine_project
+            )
             logger.info(f"Collection {self.index} exist.")
             return True
         except Exception:
@@ -101,6 +109,7 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
         client = self._get_client()
         response = client.create_collection(
             collection_name=self.index,
+            project=self.volcengine_project,
             description="Created by Volcengine Agent Development Kit VeADK",
             builtin_event_types=self.memory_type,
         )
@@ -156,7 +165,9 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
         )
 
         client = self._get_sdk_client()
-        collection = client.get_collection(collection_name=self.index)
+        collection = client.get_collection(
+            collection_name=self.index, project_name=self.volcengine_project
+        )
         response = collection.add_session(
             session_id=session_id,
             messages=messages,
@@ -181,7 +192,9 @@ class VikingDBLTMBackend(BaseLongTermMemoryBackend):
         )
 
         client = self._get_sdk_client()
-        collection = client.get_collection(collection_name=self.index)
+        collection = client.get_collection(
+            collection_name=self.index, project_name=self.volcengine_project
+        )
         response = collection.search_memory(
             query=query,
             filter=filter,

{veadk_python-0.2.35 → veadk_python-0.5.1}/veadk/memory/short_term_memory.py

@@ -32,6 +32,7 @@ from veadk.memory.short_term_memory_backends.postgresql_backend import (
 from veadk.memory.short_term_memory_backends.sqlite_backend import (
     SQLiteSTMBackend,
 )
+from veadk.models.ark_transform import build_cache_metadata
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -49,6 +50,21 @@ def wrap_get_session_with_callbacks(obj, callback_fn: Callable):
     setattr(obj, "get_session", wrapper)
 
 
+def enable_responses_api_for_session_service(result, *args, **kwargs):
+    if result and isinstance(result, Session):
+        if result.events:
+            for event in result.events:
+                if (
+                    event.actions
+                    and event.actions.state_delta
+                    and not event.cache_metadata
+                    and "response_id" in event.actions.state_delta
+                ):
+                    event.cache_metadata = build_cache_metadata(
+                        response_id=event.actions.state_delta.get("response_id"),
+                    )
+
+
 class ShortTermMemory(BaseModel):
     """Short term memory for agent execution.
 
@@ -170,6 +186,11 @@ class ShortTermMemory(BaseModel):
                         db_kwargs=self.db_kwargs, **self.backend_configs
                     ).session_service
 
+        if self.backend != "local":
+            wrap_get_session_with_callbacks(
+                self._session_service, enable_responses_api_for_session_service
+            )
+
         if self.after_load_memory_callback:
             wrap_get_session_with_callbacks(
                 self._session_service, self.after_load_memory_callback