feat(tools): add scratchpad tool for agent working memory

python/beeai_framework/tools/scratchpad/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+from beeai_framework.tools.scratchpad.scratchpad import ScratchpadInput, ScratchpadTool
+
+__all__ = ["ScratchpadInput", "ScratchpadTool"]

python/beeai_framework/tools/scratchpad/scratchpad.py

@@ -0,0 +1,343 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Agent Scratchpad Tool - Allows agents to track their reasoning and actions.
+
+This tool provides a working memory (scratchpad) where agents can:
+- Record actions they've taken
+- Store observations/results from tools
+- Review their previous reasoning
+- Avoid repeating actions
+"""
+
+import logging
+from typing import ClassVar
+
+from pydantic import BaseModel, Field
+
+from beeai_framework.context import RunContext
+from beeai_framework.emitter import Emitter
+from beeai_framework.tools import StringToolOutput, Tool, ToolRunOptions
+
+logger = logging.getLogger(__name__)
+
+
+class ScratchpadInput(BaseModel):
+    """Input schema for scratchpad operations."""
+
+    operation: str = Field(
+        description=(
+            "Operation to perform: 'read' to view scratchpad, "
+            "'write' to add entry, 'append' to add to last entry, "
+            "'clear' to reset"
+        ),
+        enum=["read", "write", "append", "clear"],
+    )
+    content: str | None = Field(
+        default=None,
+        description=(
+            "Content to write/append (required for 'write' and 'append' " "operations)"
+        ),
+    )
+
+
+class ScratchpadTool(Tool):
+    """Tool for managing agent scratchpad (working memory)."""
+
+    _scratchpads: ClassVar[dict[str, list]] = {}
+
+    def __init__(self, session_id: str | None = None) -> None:
+        """Initialize scratchpad tool.
+
+        Args:
+            session_id: Optional session identifier (deprecated, not used).
+                        Session ID is now extracted from RunContext.
+        """
+        super().__init__()
+        self.middlewares = []
+
+    @staticmethod
+    def _ensure_session(session_id: str) -> None:
+        """Ensure a session exists in scratchpads."""
+        if session_id not in ScratchpadTool._scratchpads:
+            ScratchpadTool._scratchpads[session_id] = []
+
+    def _get_session_id(self, context: RunContext | None = None) -> str:
+        """Extract session ID from context.
+
+        Always returns "default" to maintain a single persistent scratchpad
+        across all requests, ensuring information is retained between interactions.
+
+        Args:
+            context: Run context (not used, maintained for compatibility).
+
+        Returns:
+            Session ID string (always "default").
+        """
+        # Use a single persistent session for all operations
+        # This ensures the scratchpad persists across HTTP requests
+        return "default"
+
+    @property
+    def name(self) -> str:
+        """Tool name."""
+        return "scratchpad"
+
+    @property
+    def description(self) -> str:
+        """Tool description."""
+        return (
+            "Manage your working memory (scratchpad). Use this to track "
+            "what you've done, what results you got, and avoid repeating "
+            "actions. Operations: 'read' - see your scratchpad, 'write' - "
+            "add an entry, 'clear' - reset scratchpad, 'append' - add to "
+            "existing entry."
+        )
+
+    @property
+    def input_schema(self) -> type[BaseModel]:
+        """Input schema for the tool."""
+        return ScratchpadInput
+
+    def _create_emitter(self) -> Emitter:
+        """Create emitter for the tool."""
+        return Emitter()
+
+    def _get_entries(self, session_id: str) -> list:
+        """Get scratchpad entries for a session.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            List of scratchpad entries.
+        """
+        self._ensure_session(session_id)
+        return self._scratchpads[session_id]
+
+    def _read_scratchpad(self, session_id: str) -> str:
+        """Read the current scratchpad content.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            Formatted scratchpad content string.
+        """
+        entries = self._get_entries(session_id)
+        if not entries:
+            result = "Scratchpad is empty. No actions recorded yet."
+            logger.info(f"ScratchpadTool[{session_id}]: READ - Empty")
+            return result
+
+        result = "=== AGENT SCRATCHPAD ===\n\n"
+        result += "\n\n".join(f"[{i}] {entry}" for i, entry in enumerate(entries, 1))
+
+        logger.info(f"ScratchpadTool[{session_id}]: READ - {len(entries)} entries")
+        return result
+
+    @staticmethod
+    def _parse_key_value_pairs(content: str) -> dict:
+        """Parse key-value pairs from scratchpad content.
+
+        Handles formats like:
+        - "key: value"
+        - "key1: value1, key2: value2"
+        - "key: value, key2: value2, key3: value3"
+
+        Args:
+            content: Content string to parse.
+
+        Returns:
+            Dictionary of key-value pairs.
+        """
+        pairs = {}
+        # Split by comma, but be careful with commas inside values
+        parts = [p.strip() for p in content.split(",")]
+        for part in parts:
+            if ":" in part:
+                key, value = part.split(":", 1)
+                key = key.strip()
+                value = value.strip()
+                if key and value:
+                    pairs[key] = value
+        return pairs
+
+    @staticmethod
+    def _merge_entries(entries: list, new_pairs: dict) -> list:
+        """Merge new key-value pairs into existing entries.
+
+        Args:
+            entries: List of existing scratchpad entries.
+            new_pairs: Dictionary of new key-value pairs to merge.
+
+        Returns:
+            Updated list of entries (consolidated).
+        """
+        # Parse all existing entries into a single dict
+        consolidated = {}
+        for entry in entries:
+            pairs = ScratchpadTool._parse_key_value_pairs(entry)
+            consolidated.update(pairs)
+
+        # Merge new pairs (new values override old ones)
+        consolidated.update(new_pairs)
+
+        # Convert back to entry format
+        if consolidated:
+            # Create a single consolidated entry
+            entry_str = ", ".join(f"{k}: {v}" for k, v in consolidated.items())
+            return [entry_str]
+        return []
+
+    def _write_scratchpad(self, entry: str, session_id: str) -> str:
+        """Add or update entry in the scratchpad.
+
+        Merges key-value pairs with existing entries to avoid duplicates.
+        If entry contains key-value pairs (format: "key: value"), it will
+        update existing entries with the same keys.
+
+        Args:
+            entry: Content to add/update.
+            session_id: Session identifier.
+
+        Returns:
+            Success message.
+        """
+        entries = self._get_entries(session_id)
+        new_pairs = self._parse_key_value_pairs(entry)
+
+        if new_pairs:
+            # Merge with existing entries
+            entries[:] = self._merge_entries(entries, new_pairs)
+            result = f"Updated scratchpad: {', '.join(f'{k}: {v}' for k, v in new_pairs.items())}"
+        else:
+            # If no key-value pairs found, append as-is (for non-structured entries)
+            entries.append(entry)
+            result = f"Added to scratchpad: {entry}"
+
+        logger.info(
+            f"ScratchpadTool[{session_id}]: WRITE - " f"{len(entries)} total entries"
+        )
+        return result
+
+    def _append_scratchpad(self, text: str, session_id: str) -> str:
+        """Append to the last entry in scratchpad.
+
+        Args:
+            text: Text to append.
+            session_id: Session identifier.
+
+        Returns:
+            Success or error message.
+        """
+        entries = self._get_entries(session_id)
+        if not entries:
+            result = "No entry to append to. Use 'write' first."
+            logger.info(f"ScratchpadTool[{session_id}]: APPEND - No entries")
+            return result
+
+        entries[-1] += f" {text}"
+        result = f"Appended to last entry: {text}"
+        logger.info(f"ScratchpadTool[{session_id}]: APPEND - Updated")
+        return result
+
+    def _clear_scratchpad(self, session_id: str) -> str:
+        """Clear the scratchpad.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            Success message.
+        """
+        entries_count = len(self._get_entries(session_id))
+        self._scratchpads[session_id] = []
+        result = "Scratchpad cleared."
+        logger.info(
+            f"ScratchpadTool[{session_id}]: CLEAR - " f"{entries_count} entries"
+        )
+        return result
+
+    async def _run(
+        self,
+        input: ScratchpadInput,
+        options: ToolRunOptions | None = None,
+        context: RunContext | None = None,
+    ) -> StringToolOutput:
+        """Execute scratchpad operation.
+
+        Args:
+            input: ScratchpadInput model instance.
+            options: Optional tool run options.
+            context: Optional run context.
+
+        Returns:
+            StringToolOutput with the result of the operation.
+        """
+        # Get session ID (always "default" for persistent storage)
+        session_id = self._get_session_id(context)
+        operation = input.operation.lower().strip()
+        content = input.content
+
+        logger.info(
+            f"ScratchpadTool[{session_id}]: operation='{operation}', "
+            f"content='{content}'"
+        )
+
+        if not operation:
+            error_msg = (
+                "Error: 'operation' parameter is required. "
+                "Use 'read', 'write', 'append', or 'clear'."
+            )
+            return StringToolOutput(result=error_msg)
+
+        # Operation handlers
+        handlers = {
+            "read": lambda: self._read_scratchpad(session_id),
+            "write": lambda: (
+                self._write_scratchpad(content, session_id)
+                if content
+                else "Error: 'write' operation requires 'content' parameter."
+            ),
+            "append": lambda: (
+                self._append_scratchpad(content, session_id)
+                if content
+                else "Error: 'append' operation requires 'content' parameter."
+            ),
+            "clear": lambda: self._clear_scratchpad(session_id),
+        }
+
+        handler = handlers.get(operation)
+        if handler:
+            result = handler()
+            return StringToolOutput(result=result)
+
+        error_msg = (
+            f"Unknown operation: {operation}. "
+            "Use 'read', 'write', 'append', or 'clear'."
+        )
+        return StringToolOutput(result=error_msg)
+
+    @classmethod
+    def get_scratchpad_for_session(cls, session_id: str) -> list:
+        """Get scratchpad entries for a specific session.
+
+        Args:
+            session_id: Session identifier.
+
+        Returns:
+            List of scratchpad entries.
+        """
+        return cls._scratchpads.get(session_id, [])
+
+    @classmethod
+    def clear_session(cls, session_id: str) -> None:
+        """Clear scratchpad for a specific session.
+
+        Args:
+            session_id: Session identifier.
+        """
+        if session_id in cls._scratchpads:
+            cls._scratchpads[session_id] = []