Source code for aidputils.agents.toolkit.memory_helper

from __future__ import annotations

import _collections_abc
import typing
from abc import ABC, abstractmethod
from enum import StrEnum
from typing import Any, Dict, List, Optional, Final
import logging

from langgraph.checkpoint.memory import InMemorySaver

logger = logging.getLogger('memory_helper')


[docs] class AbstractCheckpointSaver(ABC): """Local abstract interface for Oracle Memory checkpointers. This is *not* LangGraph's `BaseCheckpointSaver`; it is a project-level interface that combines the public surface area of: - ProxyCheckpointClient (sync) - AsyncProxyCheckpointClient (async) It exists so code in this repo can depend on a single type that supports either sync or async execution. """ # ----- common / misc ----- # ----- sync API (ProxyCheckpointClient) -----
[docs] def health(self) -> Dict[str, Any]: raise NotImplementedError
[docs] def get_tuple(self, config: Dict[str, Any], kwargs: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]: raise NotImplementedError
[docs] def list( self, config: Optional[Dict[str, Any]], filter: Optional[Dict[str, Any]] = None, before: Optional[Dict[str, Any]] = None, limit: Optional[int] = None, kwargs: Optional[Dict[str, Any]] = None, ) -> List[Dict[str, Any]]: raise NotImplementedError
[docs] def put( self, config: Dict[str, Any], checkpoint: Dict[str, Any], metadata: Dict[str, Any], new_versions: Dict[str, Any], kwargs: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: raise NotImplementedError
[docs] def put_writes(self, *args, **options): raise NotImplementedError
[docs] def delete_thread(self, thread_id: str) -> Dict[str, Any]: raise NotImplementedError
[docs] def delete(self, config_or_thread: Any) -> Dict[str, Any]: raise NotImplementedError
# ----- async API (AsyncProxyCheckpointClient) -----
[docs] async def ahealth(self) -> Dict[str, Any]: raise NotImplementedError
[docs] async def aget_tuple(self, config: Dict[str, Any], kwargs: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]: raise NotImplementedError
[docs] async def alist( self, config: Optional[Dict[str, Any]], filter: Optional[Dict[str, Any]] = None, before: Optional[Dict[str, Any]] = None, limit: Optional[int] = None, kwargs: Optional[Dict[str, Any]] = None, ) -> List[Dict[str, Any]]: raise NotImplementedError
[docs] async def aput( self, config: Dict[str, Any], checkpoint: Dict[str, Any], metadata: Dict[str, Any], new_versions: Dict[str, Any], kwargs: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: raise NotImplementedError
[docs] async def aput_writes(self, *args, **options): raise NotImplementedError
[docs] async def adelete_thread(self, thread_id: str) -> Dict[str, Any]: raise NotImplementedError
[docs] async def adelete(self, config_or_thread: Any) -> Dict[str, Any]: raise NotImplementedError
[docs] async def aclose(self) -> None: raise NotImplementedError
[docs] class InMemoryCheckpointSaver(InMemorySaver, AbstractCheckpointSaver): # type: ignore[misc] """In-memory checkpointer that satisfies our project-level AbstractCheckpointSaver. This is mainly useful for local/dev/test runs where you want LangGraph-compatible checkpointing without a backing memory server. """ def __init__(self, *args, **kwargs): if InMemorySaver is None: # pragma: no cover raise RuntimeError("langgraph.checkpoint.memory.InMemorySaver is not available") super().__init__(*args, **kwargs) # AbstractCheckpointSaver adds convenience methods not required by LangGraph.
[docs] def health(self) -> Dict[str, Any]: return {"status": "ok", "backend": "in_memory"}
[docs] async def ahealth(self) -> Dict[str, Any]: return {"status": "ok", "backend": "in_memory"}
[docs] def delete(self, config_or_thread: Any) -> Dict[str, Any]: """Delete a thread by thread_id or config. Mirrors the convenience helpers on ProxyCheckpointClient / AsyncProxyCheckpointClient. """ thread_id = None if isinstance(config_or_thread, str): thread_id = config_or_thread elif isinstance(config_or_thread, (dict, _collections_abc.Mapping)): c = config_or_thread cfg = c.get("configurable") if isinstance(c.get("configurable"), (dict, _collections_abc.Mapping)) else None thread_id = (cfg or {}).get("thread_id") or c.get("thread_id") if not thread_id: raise ValueError("delete requires thread_id or config with configurable.thread_id") # InMemorySaver.delete_thread returns None self.delete_thread(thread_id) return {"deleted": True, "thread_id": thread_id}
[docs] async def adelete(self, config_or_thread: Any) -> Dict[str, Any]: thread_id = None if isinstance(config_or_thread, str): thread_id = config_or_thread elif isinstance(config_or_thread, (dict, _collections_abc.Mapping)): c = config_or_thread cfg = c.get("configurable") if isinstance(c.get("configurable"), (dict, _collections_abc.Mapping)) else None thread_id = (cfg or {}).get("thread_id") or c.get("thread_id") if not thread_id: raise ValueError("adelete requires thread_id or config with configurable.thread_id") await self.adelete_thread(thread_id) return {"deleted": True, "thread_id": thread_id}
[docs] async def aclose(self) -> None: # Nothing to close for in-memory saver return None
[docs] class CheckpointerType(StrEnum) : SYNC = "SYNC" ASYNC = "ASYNC" IN_MEMORY = "IN_MEMORY"
[docs] def create_checkpoint_saver(agent_name: str, platform_name: str = 'langgraph', checkpointer_type: str = 'ASYNC', base_url: str = 'http://localhost:21100/') -> Optional[AbstractCheckpointSaver]: if globals().get('checkpointer_v1') is not None: return globals().get('checkpointer_v1') checkpointer = None checkpointer_type = checkpointer_type.upper() if "langgraph" == platform_name.lower(): if checkpointer_type == CheckpointerType.SYNC.value : try: from oracle_memory_clients.client import ProxyCheckpointClient checkpointer = ProxyCheckpointClient(agent=agent_name, base_url=base_url) except ImportError: logger.warn("Memory server is not available (AsyncProxyCheckpointClient class not found), " "no checkpointer will be used for this session.") elif checkpointer_type == CheckpointerType.ASYNC.value: try: from oracle_memory_clients.client import AsyncProxyCheckpointClient checkpointer = AsyncProxyCheckpointClient(agent=agent_name, base_url=base_url) except ImportError: logger.warn("Memory server is not available (AsyncProxyCheckpointClient class not found), " "no checkpointer will be used for this session.") elif checkpointer_type == CheckpointerType.IN_MEMORY.value: checkpointer = InMemoryCheckpointSaver(agent=agent_name, base_url=base_url) elif checkpointer_type is None: logger.info(f"Checkpointer type is specified as None for {agent_name}") else: raise ValueError("checkpointer_type must be one of 'SYNC', 'ASYNC', 'IN_MEMORY', or None") if checkpointer is not None: logger.info(f"Created checkpointer for {agent_name}, type: {checkpointer_type}") globals().setdefault('checkpointer_v1', checkpointer) else: logger.info(f"No checkpointer was created for {agent_name}.") return checkpointer