|
29 | 29 | ) |
30 | 30 |
|
31 | 31 | from .bedrock_converter import AgentCoreMemoryConverter |
32 | | -from .config import AgentCoreMemoryConfig, RetrievalConfig, normalize_metadata |
| 32 | +from .config import AgentCoreMemoryConfig, PersistenceMode, RetrievalConfig, normalize_metadata |
33 | 33 | from .converters import MemoryConverter |
34 | 34 |
|
35 | 35 | if TYPE_CHECKING: |
@@ -142,6 +142,7 @@ def __init__( |
142 | 142 | """ |
143 | 143 | self.converter = converter or AgentCoreMemoryConverter |
144 | 144 | self.config = agentcore_memory_config |
| 145 | + self.persistence_mode = agentcore_memory_config.persistence_mode |
145 | 146 | self.memory_client = MemoryClient(region_name=region_name) |
146 | 147 | session = boto_session or boto3.Session(region_name=region_name) |
147 | 148 | self.has_existing_agent = False |
@@ -255,17 +256,19 @@ def create_session(self, session: Session, **kwargs: Any) -> Session: |
255 | 256 | if session.session_id != self.config.session_id: |
256 | 257 | raise SessionException(f"Session ID mismatch: expected {self.config.session_id}, got {session.session_id}") |
257 | 258 |
|
258 | | - event = self.memory_client.gmdp_client.create_event( |
259 | | - memoryId=self.config.memory_id, |
260 | | - actorId=self.config.actor_id, |
261 | | - sessionId=self.session_id, |
262 | | - payload=[ |
263 | | - {"blob": json.dumps(session.to_dict())}, |
264 | | - ], |
265 | | - eventTimestamp=self._get_monotonic_timestamp(), |
266 | | - metadata={STATE_TYPE_KEY: {"stringValue": StateType.SESSION.value}}, |
267 | | - ) |
268 | | - logger.info("Created session: %s with event: %s", session.session_id, event.get("event", {}).get("eventId")) |
| 259 | + if self.persistence_mode is not PersistenceMode.NONE: |
| 260 | + event = self.memory_client.gmdp_client.create_event( |
| 261 | + memoryId=self.config.memory_id, |
| 262 | + actorId=self.config.actor_id, |
| 263 | + sessionId=self.session_id, |
| 264 | + payload=[ |
| 265 | + {"blob": json.dumps(session.to_dict())}, |
| 266 | + ], |
| 267 | + eventTimestamp=self._get_monotonic_timestamp(), |
| 268 | + metadata={STATE_TYPE_KEY: {"stringValue": StateType.SESSION.value}}, |
| 269 | + ) |
| 270 | + logger.info("Created session: %s with event: %s", session.session_id, event.get("event", {}).get("eventId")) |
| 271 | + |
269 | 272 | return session |
270 | 273 |
|
271 | 274 | def read_session(self, session_id: str, **kwargs: Any) -> Optional[Session]: |
@@ -318,14 +321,15 @@ def read_session(self, session_id: str, **kwargs: Any) -> Optional[Session]: |
318 | 321 | session_data = json.loads(old_event.get("payload", {})[0].get("blob")) |
319 | 322 | session = Session.from_dict(session_data) |
320 | 323 | # Migrate: create new event with metadata, delete old |
321 | | - self.create_session(session) |
322 | | - self.memory_client.gmdp_client.delete_event( |
323 | | - memoryId=self.config.memory_id, |
324 | | - actorId=legacy_actor_id, |
325 | | - sessionId=session_id, |
326 | | - eventId=old_event.get("eventId"), |
327 | | - ) |
328 | | - logger.info("Migrated legacy session event for session: %s", session_id) |
| 324 | + if self.persistence_mode is not PersistenceMode.NONE: |
| 325 | + self.create_session(session) |
| 326 | + self.memory_client.gmdp_client.delete_event( |
| 327 | + memoryId=self.config.memory_id, |
| 328 | + actorId=legacy_actor_id, |
| 329 | + sessionId=session_id, |
| 330 | + eventId=old_event.get("eventId"), |
| 331 | + ) |
| 332 | + logger.info("Migrated legacy session event for session: %s", session_id) |
329 | 333 | return session |
330 | 334 |
|
331 | 335 | return None |
@@ -364,6 +368,9 @@ def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: A |
364 | 368 | if session_agent.created_at: |
365 | 369 | self._agent_created_at_cache[session_agent.agent_id] = session_agent.created_at |
366 | 370 |
|
| 371 | + if self.persistence_mode is PersistenceMode.NONE: |
| 372 | + return |
| 373 | + |
367 | 374 | if self.config.batch_size > 1: |
368 | 375 | # Buffer the agent state events |
369 | 376 | should_flush = False |
@@ -462,14 +469,15 @@ def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[ |
462 | 469 | agent_data = json.loads(old_event.get("payload", {})[0].get("blob")) |
463 | 470 | agent = SessionAgent.from_dict(agent_data) |
464 | 471 | # Migrate: create new event with metadata, delete old |
465 | | - self.create_agent(session_id, agent) |
466 | | - self.memory_client.gmdp_client.delete_event( |
467 | | - memoryId=self.config.memory_id, |
468 | | - actorId=legacy_actor_id, |
469 | | - sessionId=session_id, |
470 | | - eventId=old_event.get("eventId"), |
471 | | - ) |
472 | | - logger.info("Migrated legacy agent event for agent: %s", agent_id) |
| 472 | + if self.persistence_mode is not PersistenceMode.NONE: |
| 473 | + self.create_agent(session_id, agent) |
| 474 | + self.memory_client.gmdp_client.delete_event( |
| 475 | + memoryId=self.config.memory_id, |
| 476 | + actorId=legacy_actor_id, |
| 477 | + sessionId=session_id, |
| 478 | + eventId=old_event.get("eventId"), |
| 479 | + ) |
| 480 | + logger.info("Migrated legacy agent event for agent: %s", agent_id) |
473 | 481 | return agent |
474 | 482 |
|
475 | 483 | return None |
@@ -546,6 +554,9 @@ def create_message( |
546 | 554 | if not messages: |
547 | 555 | return None |
548 | 556 |
|
| 557 | + if self.persistence_mode is PersistenceMode.NONE: |
| 558 | + return {} |
| 559 | + |
549 | 560 | is_blob = self.converter.exceeds_conversational_limit(messages[0]) |
550 | 561 |
|
551 | 562 | # Build merged metadata from config defaults + per-call overrides |
@@ -862,6 +873,9 @@ def _flush_messages_only(self) -> list[dict[str, Any]]: |
862 | 873 | Raises: |
863 | 874 | SessionException: If message creation fails. On failure, messages remain in the buffer. |
864 | 875 | """ |
| 876 | + if self.persistence_mode is PersistenceMode.NONE: |
| 877 | + return [] |
| 878 | + |
865 | 879 | with self._message_lock: |
866 | 880 | messages_to_send = list(self._message_buffer) |
867 | 881 |
|
@@ -940,6 +954,9 @@ def _flush_agent_states_only(self) -> list[dict[str, Any]]: |
940 | 954 | Raises: |
941 | 955 | SessionException: If agent state creation fails. On failure, agent states remain in the buffer. |
942 | 956 | """ |
| 957 | + if self.persistence_mode is PersistenceMode.NONE: |
| 958 | + return [] |
| 959 | + |
943 | 960 | with self._agent_state_lock: |
944 | 961 | agent_states_to_send = list(self._agent_state_buffer) |
945 | 962 |
|
|
0 commit comments