aip-agents-binary 0.5.21__py3-none-macosx_13_0_arm64.whl → 0.6.8__py3-none-macosx_13_0_arm64.whl

aip_agents/tools/gl_connector_tools.py

@@ -0,0 +1,119 @@
+"""Auto-generated tools from GL Connectors.
+
+Authors:
+    Saul Sayers (saul.sayers@gdplabs.id)
+"""
+
+from bosa_connectors import BosaConnector, BOSAConnectorToolGenerator
+from langchain_core.tools import BaseTool
+
+from aip_agents.tools.constants import (
+    GL_CONNECTORS_API_KEY,
+    GL_CONNECTORS_BASE_URL,
+    GL_CONNECTORS_FETCH_MAX_RETRIES,
+    ToolType,
+)
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_gl_connector_modules_with_retry() -> list[str]:
+    """Try to get available modules with retries.
+
+    Returns:
+        List of available modules.
+    """
+    if not GL_CONNECTORS_BASE_URL or not GL_CONNECTORS_API_KEY:
+        logger.warning("GL Connectors credentials missing (base_url or api_key); returning empty modules list")
+        return []
+
+    connector = BosaConnector(api_base_url=GL_CONNECTORS_BASE_URL, api_key=GL_CONNECTORS_API_KEY)
+    modules = []
+    for attempt in range(GL_CONNECTORS_FETCH_MAX_RETRIES):
+        try:
+            modules = list(connector.get_available_modules())
+            if modules:
+                return modules
+            logger.warning(
+                f"Failed to get GL Connectors available modules, retrying... (attempt {attempt + 1} / {GL_CONNECTORS_FETCH_MAX_RETRIES})"
+            )
+        except Exception as e:
+            logger.exception(
+                f"Exception when getting GL Connectors available modules (attempt {attempt + 1} / {GL_CONNECTORS_FETCH_MAX_RETRIES}): {e}"
+            )
+    logger.error("Failed to get GL Connectors available modules after maximum retries")
+    return modules
+
+
+class LazyGLConnectorToolsDict(dict):
+    """Lazy dictionary for GL Connectors."""
+
+    def __missing__(self, app):
+        """When a key is missing, create the tools and store them in the dictionary.
+
+        Args:
+            app: Name of the GL Connectors.
+
+        Returns:
+            List of tools generated by the GL Connectors tool generator.
+        """
+        if app not in get_gl_connector_modules():
+            return []
+        tools = []
+        for attempt in range(GL_CONNECTORS_FETCH_MAX_RETRIES):
+            try:
+                tools = BOSAConnectorToolGenerator(
+                    api_base_url=GL_CONNECTORS_BASE_URL,
+                    api_key=GL_CONNECTORS_API_KEY,
+                    app_name=app,
+                ).generate_tools(tool_type=ToolType.LANGCHAIN)
+                if tools:
+                    self[app] = tools
+                    return tools
+                logger.warning(
+                    f"Failed to create GL Connectors, retrying... (attempt {attempt + 1} / {GL_CONNECTORS_FETCH_MAX_RETRIES})"
+                )
+            except Exception as e:
+                logger.exception(
+                    f"Exception when creating GL Connectors for app '{app}' "
+                    f"(attempt {attempt + 1} / {GL_CONNECTORS_FETCH_MAX_RETRIES}): {e}"
+                )
+        logger.error("Failed to create GL Connectors after maximum retries")
+        return tools
+
+
+# Supported modules (dynamic)
+def get_gl_connector_modules() -> list[str]:
+    """Lazily fetch and cache GL Connectors modules.
+
+    Returns:
+        List of GL Connectors modules.
+    """
+    if not hasattr(get_gl_connector_modules, "_cache"):
+        get_gl_connector_modules._cache = get_gl_connector_modules_with_retry()
+    return get_gl_connector_modules._cache
+
+
+# FOR BACKWARDS COMPATIBILITY
+def get_bosa_modules() -> list[str]:
+    """Backward-compatible alias for get_gl_connector_modules."""
+    return get_gl_connector_modules()
+
+
+GL_CONNECTORS_MODULES = [
+    "github",
+    "twitter",
+    "google",
+    "google_drive",
+    "google_mail",
+    "google_docs",
+]
+
+# FOR BACKWARDS COMPATIBILITY
+BOSA_MODULES = GL_CONNECTORS_MODULES
+
+GL_CONNECTORS_AUTOMATED_TOOLS: dict[str, list[BaseTool]] = LazyGLConnectorToolsDict()
+
+# FOR BACKWARDS COMPATIBILITY
+BOSA_AUTOMATED_TOOLS = GL_CONNECTORS_AUTOMATED_TOOLS

aip_agents/tools/gl_connector_tools.pyi

@@ -0,0 +1,39 @@
+from _typeshed import Incomplete
+from aip_agents.tools.constants import GL_CONNECTORS_API_KEY as GL_CONNECTORS_API_KEY, GL_CONNECTORS_BASE_URL as GL_CONNECTORS_BASE_URL, GL_CONNECTORS_FETCH_MAX_RETRIES as GL_CONNECTORS_FETCH_MAX_RETRIES, ToolType as ToolType
+from aip_agents.utils.logger import get_logger as get_logger
+from langchain_core.tools import BaseTool as BaseTool
+
+logger: Incomplete
+
+def get_gl_connector_modules_with_retry() -> list[str]:
+    """Try to get available modules with retries.
+
+    Returns:
+        List of available modules.
+    """
+
+class LazyGLConnectorToolsDict(dict):
+    """Lazy dictionary for GL Connectors."""
+    def __missing__(self, app):
+        """When a key is missing, create the tools and store them in the dictionary.
+
+        Args:
+            app: Name of the GL Connectors.
+
+        Returns:
+            List of tools generated by the GL Connectors tool generator.
+        """
+
+def get_gl_connector_modules() -> list[str]:
+    """Lazily fetch and cache GL Connectors modules.
+
+    Returns:
+        List of GL Connectors modules.
+    """
+def get_bosa_modules() -> list[str]:
+    """Backward-compatible alias for get_gl_connector_modules."""
+
+GL_CONNECTORS_MODULES: Incomplete
+BOSA_MODULES = GL_CONNECTORS_MODULES
+GL_CONNECTORS_AUTOMATED_TOOLS: dict[str, list[BaseTool]]
+BOSA_AUTOMATED_TOOLS = GL_CONNECTORS_AUTOMATED_TOOLS

aip_agents/tools/memory_search/__init__.py

@@ -6,17 +6,24 @@ Authors:
 
 from aip_agents.tools.memory_search.base import LongTermMemorySearchTool
 from aip_agents.tools.memory_search.mem0 import (
+    MEMORY_DELETE_TOOL_NAME,
     MEMORY_SEARCH_TOOL_NAME,
+    Mem0DeleteInput,
+    Mem0DeleteTool,
     Mem0SearchInput,
     Mem0SearchTool,
 )
-from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput, MemoryConfig
+from aip_agents.tools.memory_search.schema import LongTermMemoryDeleteInput, LongTermMemorySearchInput, MemoryConfig
 
 __all__ = [
     "MemoryConfig",
+    "LongTermMemoryDeleteInput",
     "LongTermMemorySearchInput",
     "LongTermMemorySearchTool",
+    "Mem0DeleteInput",
+    "Mem0DeleteTool",
     "Mem0SearchInput",
     "Mem0SearchTool",
+    "MEMORY_DELETE_TOOL_NAME",
     "MEMORY_SEARCH_TOOL_NAME",
 ]

aip_agents/tools/memory_search/__init__.pyi

@@ -1,5 +1,5 @@
 from aip_agents.tools.memory_search.base import LongTermMemorySearchTool as LongTermMemorySearchTool
-from aip_agents.tools.memory_search.mem0 import MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0SearchInput as Mem0SearchInput, Mem0SearchTool as Mem0SearchTool
-from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput as LongTermMemorySearchInput, MemoryConfig as MemoryConfig
+from aip_agents.tools.memory_search.mem0 import MEMORY_DELETE_TOOL_NAME as MEMORY_DELETE_TOOL_NAME, MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0DeleteInput as Mem0DeleteInput, Mem0DeleteTool as Mem0DeleteTool, Mem0SearchInput as Mem0SearchInput, Mem0SearchTool as Mem0SearchTool
+from aip_agents.tools.memory_search.schema import LongTermMemoryDeleteInput as LongTermMemoryDeleteInput, LongTermMemorySearchInput as LongTermMemorySearchInput, MemoryConfig as MemoryConfig
 
-__all__ = ['MemoryConfig', 'LongTermMemorySearchInput', 'LongTermMemorySearchTool', 'Mem0SearchInput', 'Mem0SearchTool', 'MEMORY_SEARCH_TOOL_NAME']
+__all__ = ['MemoryConfig', 'LongTermMemoryDeleteInput', 'LongTermMemorySearchInput', 'LongTermMemorySearchTool', 'Mem0DeleteInput', 'Mem0DeleteTool', 'Mem0SearchInput', 'Mem0SearchTool', 'MEMORY_DELETE_TOOL_NAME', 'MEMORY_SEARCH_TOOL_NAME']

aip_agents/tools/memory_search/mem0.py

@@ -13,13 +13,14 @@ from langchain_core.runnables import RunnableConfig
 
 from aip_agents.memory.constants import MemoryDefaults
 from aip_agents.tools.memory_search.base import LongTermMemorySearchTool
-from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput
+from aip_agents.tools.memory_search.schema import LongTermMemoryDeleteInput, LongTermMemorySearchInput
 from aip_agents.utils.datetime import is_valid_date_string, next_day_iso
 from aip_agents.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
 MEMORY_SEARCH_TOOL_NAME = "built_in_mem0_search"
+MEMORY_DELETE_TOOL_NAME = "built_in_mem0_delete"
 
 
 class Mem0SearchTool(LongTermMemorySearchTool):
@@ -256,3 +257,115 @@ class Mem0SearchTool(LongTermMemorySearchTool):
 
 
 Mem0SearchInput = LongTermMemorySearchInput
+
+
+class Mem0DeleteTool(LongTermMemorySearchTool):
+    """Mem0-specific implementation of the long-term memory delete tool."""
+
+    name: str = MEMORY_DELETE_TOOL_NAME
+    description: str = (
+        "Delete memories from long-term mem0 storage. Supports three modes:\n"
+        "1. DELETE BY IDS: Provide 'memory_ids'\n"
+        "2. DELETE BY QUERY: Provide 'query'\n"
+        "3. DELETE ALL: Provide 'delete_all=true' with no query/IDs\n"
+    )
+    args_schema: type[LongTermMemoryDeleteInput] = LongTermMemoryDeleteInput
+    LOG_PREFIX: ClassVar[str] = "Mem0DeleteTool"
+    METADATA_FILTER_BLOCKLIST: ClassVar[set[str]] = {"user_id", "memory_user_id"}
+
+    async def _arun(
+        self,
+        query: str | None = None,
+        config: RunnableConfig | None = None,
+        run_manager: Any | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """Execute the memory delete asynchronously for LangChain.
+
+        Args:
+            query: Semantic delete query when provided.
+            config: Runnable configuration containing LangChain metadata.
+            run_manager: LangChain callbacks (unused).
+            **kwargs: Additional arguments such as ``memory_ids``, ``delete_all``, ``metadata``.
+
+        Returns:
+            str: JSON-encoded delete result or an error message.
+        """
+        logger.info("%s: Received config: %s", self.LOG_PREFIX, config)
+
+        memory_ids: list[str] | None = kwargs.get("memory_ids")
+        delete_all: bool | None = kwargs.get("delete_all")
+        threshold: float | None = kwargs.get("threshold")
+        top_k: int | None = kwargs.get("top_k")
+        categories: list[str] | None = kwargs.get("categories")
+        metadata: dict[str, Any] | None = kwargs.get("metadata")
+
+        user_id = self._resolve_user_id(metadata=metadata, config=config)
+
+        metadata_filter = None
+        if isinstance(metadata, dict):
+            metadata_filter = {k: v for k, v in metadata.items() if k not in self.METADATA_FILTER_BLOCKLIST} or None
+
+        if memory_ids:
+            if not hasattr(self.memory, "delete"):
+                return f"Error executing memory tool '{self.name}': backend does not support delete()"
+            mode = "ids"
+            result = self.memory.delete(  # type: ignore[attr-defined]
+                memory_ids=memory_ids,
+                user_id=user_id,
+                metadata=metadata_filter,
+                categories=categories,
+            )
+        elif query:
+            if not hasattr(self.memory, "delete_by_query"):
+                return f"Error executing memory tool '{self.name}': backend does not support delete_by_query()"
+            mode = "query"
+            filters: dict[str, Any] | None = None
+            if metadata_filter or categories:
+                filters = {}
+                if metadata_filter:
+                    filters["metadata"] = metadata_filter
+                if categories:
+                    filters["categories"] = categories
+            result = self.memory.delete_by_query(  # type: ignore[attr-defined]
+                query=query,
+                user_id=user_id,
+                threshold=threshold,
+                top_k=top_k,
+                filters=filters,
+            )
+        elif delete_all:
+            if not hasattr(self.memory, "delete"):
+                return f"Error executing memory tool '{self.name}': backend does not support delete()"
+            mode = "all"
+            result = self.memory.delete(  # type: ignore[attr-defined]
+                memory_ids=None,
+                user_id=user_id,
+                metadata=metadata_filter,
+                categories=categories,
+            )
+        else:
+            return f"Error executing memory tool '{self.name}': provide memory_ids, query, or delete_all=true."
+
+        count = None
+        if isinstance(result, dict):
+            count = result.get("count") or result.get("deleted") or result.get("total")
+
+        logger.info(
+            "%s: delete mode=%s user_id='%s' count=%s",
+            self.LOG_PREFIX,
+            mode,
+            user_id,
+            count if count is not None else "unknown",
+        )
+
+        payload = {"status": "success", "mode": mode}
+        try:
+            json.dumps(result)
+            payload["result"] = result
+        except TypeError:
+            payload["result"] = str(result)
+        return json.dumps(payload)
+
+
+Mem0DeleteInput = LongTermMemoryDeleteInput

aip_agents/tools/memory_search/mem0.pyi

@@ -1,13 +1,14 @@
 from _typeshed import Incomplete
 from aip_agents.memory.constants import MemoryDefaults as MemoryDefaults
 from aip_agents.tools.memory_search.base import LongTermMemorySearchTool as LongTermMemorySearchTool
-from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput as LongTermMemorySearchInput
+from aip_agents.tools.memory_search.schema import LongTermMemoryDeleteInput as LongTermMemoryDeleteInput, LongTermMemorySearchInput as LongTermMemorySearchInput
 from aip_agents.utils.datetime import is_valid_date_string as is_valid_date_string, next_day_iso as next_day_iso
 from aip_agents.utils.logger import get_logger as get_logger
 from typing import ClassVar
 
 logger: Incomplete
 MEMORY_SEARCH_TOOL_NAME: str
+MEMORY_DELETE_TOOL_NAME: str
 
 class Mem0SearchTool(LongTermMemorySearchTool):
     """Mem0-specific implementation of the long-term memory search tool."""
@@ -17,3 +18,12 @@ class Mem0SearchTool(LongTermMemorySearchTool):
     LOG_PREFIX: ClassVar[str]
     METADATA_FILTER_BLOCKLIST: ClassVar[set[str]]
 Mem0SearchInput = LongTermMemorySearchInput
+
+class Mem0DeleteTool(LongTermMemorySearchTool):
+    """Mem0-specific implementation of the long-term memory delete tool."""
+    name: str
+    description: str
+    args_schema: type[LongTermMemoryDeleteInput]
+    LOG_PREFIX: ClassVar[str]
+    METADATA_FILTER_BLOCKLIST: ClassVar[set[str]]
+Mem0DeleteInput = LongTermMemoryDeleteInput

aip_agents/tools/memory_search/schema.py

@@ -46,3 +46,36 @@ class LongTermMemorySearchInput(BaseModel):
         None,
         description="Optional metadata dict to filter by (exact key-value match).",
     )
+
+
+class LongTermMemoryDeleteInput(BaseModel):
+    """Input schema for unified long-term memory deletion."""
+
+    query: str | None = Field(
+        None,
+        description="Semantic query describing memories to delete. If provided, delete_by_user_query is used.",
+    )
+    memory_ids: list[str] | None = Field(
+        None,
+        description="Optional list of memory IDs to delete directly.",
+    )
+    delete_all: bool | None = Field(
+        None,
+        description="When True and no query/IDs are provided, delete all memories for the user scope.",
+    )
+    top_k: int | None = Field(
+        None,
+        description="Optional maximum number of memories to delete by query.",
+    )
+    threshold: float | None = Field(
+        None,
+        description="Optional semantic threshold for delete_by_user_query.",
+    )
+    categories: list[str] | None = Field(
+        None,
+        description="Optional categories to filter by (uses 'in' operator).",
+    )
+    metadata: dict[str, Any] | None = Field(
+        None,
+        description="Optional metadata dict to filter by (exact key-value match).",
+    )

aip_agents/tools/memory_search/schema.pyi

@@ -13,3 +13,13 @@ class LongTermMemorySearchInput(BaseModel):
     limit: int | None
     categories: list[str] | None
     metadata: dict[str, Any] | None
+
+class LongTermMemoryDeleteInput(BaseModel):
+    """Input schema for unified long-term memory deletion."""
+    query: str | None
+    memory_ids: list[str] | None
+    delete_all: bool | None
+    top_k: int | None
+    threshold: float | None
+    categories: list[str] | None
+    metadata: dict[str, Any] | None

aip_agents/tools/memory_search_tool.py

@@ -8,9 +8,13 @@ Authors:
 """
 
 from aip_agents.tools.memory_search import (
+    MEMORY_DELETE_TOOL_NAME,
     MEMORY_SEARCH_TOOL_NAME,
+    LongTermMemoryDeleteInput,
     LongTermMemorySearchInput,
     LongTermMemorySearchTool,
+    Mem0DeleteInput,
+    Mem0DeleteTool,
     Mem0SearchInput,
     Mem0SearchTool,
     MemoryConfig,
@@ -18,9 +22,13 @@ from aip_agents.tools.memory_search import (
 
 __all__ = [
     "MemoryConfig",
+    "LongTermMemoryDeleteInput",
     "LongTermMemorySearchInput",
     "LongTermMemorySearchTool",
+    "Mem0DeleteInput",
+    "Mem0DeleteTool",
     "Mem0SearchInput",
     "Mem0SearchTool",
+    "MEMORY_DELETE_TOOL_NAME",
     "MEMORY_SEARCH_TOOL_NAME",
 ]

aip_agents/tools/memory_search_tool.pyi

@@ -1,3 +1,3 @@
-from aip_agents.tools.memory_search import LongTermMemorySearchInput as LongTermMemorySearchInput, LongTermMemorySearchTool as LongTermMemorySearchTool, MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0SearchInput as Mem0SearchInput, Mem0SearchTool as Mem0SearchTool, MemoryConfig as MemoryConfig
+from aip_agents.tools.memory_search import LongTermMemoryDeleteInput as LongTermMemoryDeleteInput, LongTermMemorySearchInput as LongTermMemorySearchInput, LongTermMemorySearchTool as LongTermMemorySearchTool, MEMORY_DELETE_TOOL_NAME as MEMORY_DELETE_TOOL_NAME, MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0DeleteInput as Mem0DeleteInput, Mem0DeleteTool as Mem0DeleteTool, Mem0SearchInput as Mem0SearchInput, Mem0SearchTool as Mem0SearchTool, MemoryConfig as MemoryConfig
 
-__all__ = ['MemoryConfig', 'LongTermMemorySearchInput', 'LongTermMemorySearchTool', 'Mem0SearchInput', 'Mem0SearchTool', 'MEMORY_SEARCH_TOOL_NAME']
+__all__ = ['MemoryConfig', 'LongTermMemoryDeleteInput', 'LongTermMemorySearchInput', 'LongTermMemorySearchTool', 'Mem0DeleteInput', 'Mem0DeleteTool', 'Mem0SearchInput', 'Mem0SearchTool', 'MEMORY_DELETE_TOOL_NAME', 'MEMORY_SEARCH_TOOL_NAME']

aip_agents/utils/langgraph/tool_managers/delegation_tool_manager.py

@@ -134,7 +134,19 @@ class DelegationToolManager(BaseLangGraphToolManager):
             Returns:
                 The result from the delegated agent, including artifacts if any.
             """
-            writer: StreamWriter = get_stream_writer()
+            try:
+                writer: StreamWriter = get_stream_writer()
+            except Exception as exc:
+                logger.warning(
+                    "DelegationToolManager: Stream writer unavailable; delegation streaming disabled.",
+                    extra={"error": str(exc), "error_type": type(exc).__name__},
+                )
+
+                def _noop_writer(_: Any) -> None:
+                    """No-op writer for non-graph execution contexts."""
+                    return None
+
+                writer = _noop_writer
 
             try:
                 # Check delegation depth limit before executing
@@ -295,6 +307,19 @@ class DelegationToolManager(BaseLangGraphToolManager):
             cfg_conf = cfg.get("configurable") if isinstance(cfg, dict) else None
             if isinstance(cfg_conf, dict):
                 parent_step_id = cfg_conf.get("parent_step_id")
+                if not parent_step_id:
+                    metadata = cfg.get("metadata") or {}
+                    tool_call_id = (
+                        cfg_conf.get("tool_call_id")
+                        or metadata.get("tool_call_id")
+                        or metadata.get("id")
+                        or (metadata.get("tool_call") or {}).get("id")
+                    )
+                    if tool_call_id and self.parent_agent is not None:
+                        thread_key = getattr(self.parent_agent, "thread_id_key", "thread_id")
+                        thread_id = cfg_conf.get(thread_key)
+                        parent_map = self.parent_agent._tool_parent_map_by_thread.get(str(thread_id), {})
+                        parent_step_id = parent_map.get(str(tool_call_id))
             _DELEGATION_PARENT_STEP_ID_CVAR.set(parent_step_id)
         except Exception:
             _DELEGATION_PARENT_STEP_ID_CVAR.set(None)

aip_agents/utils/langgraph/tool_output_management.py

@@ -471,6 +471,86 @@ class ToolOutputManager:
         """
         return self._generate_json_summary(thread_id, max_entries)
 
+    def get_latest_reference(self, thread_id: str) -> str | None:
+        """Return the most recent tool output reference for a thread.
+
+        Args:
+            thread_id: Thread ID to retrieve the latest output reference for.
+
+        Returns:
+            Latest tool output reference string or None when unavailable.
+        """
+        try:
+            summary = json.loads(self.generate_summary(thread_id, max_entries=1))
+        except Exception as exc:
+            logger.debug("Failed to parse tool output summary: %s", exc)
+            return None
+
+        if not summary:
+            return None
+        latest = summary[0].get("reference")
+        return latest if isinstance(latest, str) and latest else None
+
+    def has_reference(self, value: Any) -> bool:
+        """Check whether a value contains a tool output reference.
+
+        Args:
+            value: Value to inspect for tool output references.
+
+        Returns:
+            True if any tool output reference is present.
+        """
+        if isinstance(value, str):
+            return value.startswith(TOOL_OUTPUT_REFERENCE_PREFIX)
+        if isinstance(value, dict):
+            return any(self.has_reference(item) for item in value.values())
+        if isinstance(value, list):
+            return any(self.has_reference(item) for item in value)
+        return False
+
+    def should_replace_with_reference(self, value: Any) -> bool:
+        """Check whether a tool argument value should use a tool output reference.
+
+        Args:
+            value: Value to evaluate for replacement.
+
+        Returns:
+            True if the value should be replaced with a reference.
+        """
+        if isinstance(value, dict | list | tuple):
+            return True
+        if isinstance(value, str):
+            return len(value) > DATA_PREVIEW_TRUNCATION_LENGTH
+        return False
+
+    def rewrite_args_with_latest_reference(self, args: dict[str, Any], thread_id: str) -> dict[str, Any]:
+        """Rewrite tool args to use the latest tool output reference when appropriate.
+
+        Args:
+            args: Tool arguments to rewrite.
+            thread_id: Thread ID used for resolving stored outputs.
+
+        Returns:
+            Updated args dictionary with references substituted when needed.
+        """
+        if not self.has_outputs(thread_id):
+            return args
+        if self.has_reference(args):
+            return args
+
+        latest_reference = self.get_latest_reference(thread_id)
+        if not latest_reference:
+            return args
+
+        updated_args = dict(args)
+        replaced_any = False
+        for key, value in args.items():
+            if self.should_replace_with_reference(value):
+                updated_args[key] = latest_reference
+                replaced_any = True
+
+        return updated_args if replaced_any else args
+
     def _generate_json_summary(self, thread_id: str, max_entries: int) -> str:
         """Generate simplified JSON summary optimized for LLM prompts.
 

aip_agents/utils/langgraph/tool_output_management.pyi

@@ -232,6 +232,43 @@ class ToolOutputManager:
             A JSON string containing structured data about tool outputs. Always returns
             valid JSON, even when no outputs are stored (empty entries list).
         """
+    def get_latest_reference(self, thread_id: str) -> str | None:
+        """Return the most recent tool output reference for a thread.
+
+        Args:
+            thread_id: Thread ID to retrieve the latest output reference for.
+
+        Returns:
+            Latest tool output reference string or None when unavailable.
+        """
+    def has_reference(self, value: Any) -> bool:
+        """Check whether a value contains a tool output reference.
+
+        Args:
+            value: Value to inspect for tool output references.
+
+        Returns:
+            True if any tool output reference is present.
+        """
+    def should_replace_with_reference(self, value: Any) -> bool:
+        """Check whether a tool argument value should use a tool output reference.
+
+        Args:
+            value: Value to evaluate for replacement.
+
+        Returns:
+            True if the value should be replaced with a reference.
+        """
+    def rewrite_args_with_latest_reference(self, args: dict[str, Any], thread_id: str) -> dict[str, Any]:
+        """Rewrite tool args to use the latest tool output reference when appropriate.
+
+        Args:
+            args: Tool arguments to rewrite.
+            thread_id: Thread ID used for resolving stored outputs.
+
+        Returns:
+            Updated args dictionary with references substituted when needed.
+        """
     def clear_all(self) -> None:
         """Clear all stored outputs from both metadata and storage.