railtracks.prebuilt.tools

 32class KeyValueMemoryToolSet(ToolSet):
 33    """Prebuilt key-value memory tools for an agent.
 34
 35    Gives an agent a persistent, exact-match scratch pad: save a fact under a
 36    key, read it back later, forget it, list everything, or search. State lives
 37    in the injected :class:`~railtracks.retrieval.stores.key_value.KeyValueStore`
 38    (defaults to an in-process :class:`InMemoryKeyValueStore`). Pass a store
 39    constructed with a ``snapshot_path`` for persistence across runs::
 40
 41        store = InMemoryKeyValueStore(snapshot_path="memory.json")
 42        toolset = KeyValueMemoryToolSet(store=store)
 43
 44    All memory in a toolset shares one namespace. To keep separate memories for
 45    different agents, give each its own ``KeyValueMemoryToolSet`` (and its own
 46    store).
 47
 48    Args:
 49        store: Backing key-value store. Defaults to a fresh, ephemeral
 50            ``InMemoryKeyValueStore``.
 51        search: Ranking algorithm used by ``search_memories``. Defaults to
 52            ``LexicalSearch()``. Pass a ``LexicalSearch(LexicalSearchConfig(...))``
 53            to tune the ranking weights, pass a
 54            ``SemanticSearch(embedding=...)`` for dense-vector ranking, or any
 55            other :class:`~railtracks.prebuilt.tools.memory.search.SearchAlgorithm`
 56            implementation to swap the algorithm entirely.
 57        on_change: Optional callback fired after every mutation, letting an
 58            outer system react (push to a UI, mirror to a database, log).
 59            Called as ``on_change(key, value)`` where ``value`` is the new
 60            value on a save and ``None`` on a forget. Exceptions raised by the
 61            callback are logged and swallowed so they never break a tool call.
 62    """
 63
 64    def __init__(
 65        self,
 66        store: KeyValueStore | None = None,
 67        search: SearchAlgorithm | None = None,
 68        on_change: Callable[[str, str | None], None] | None = None,
 69    ) -> None:
 70        self.store: KeyValueStore = store if store is not None else _default_store()
 71        self.search: SearchAlgorithm = (
 72            search if search is not None else _default_search()
 73        )
 74        self.on_change = on_change
 75
 76    def _notify(self, key: str, value: str | None) -> None:
 77        if self.on_change is None:
 78            return
 79        try:
 80            self.on_change(key, value)
 81        except Exception as e:
 82            logger.error(f"Error in on_change callback for key {key!r}: {e}")
 83
 84    async def remember(self, key: str, value: str) -> str:
 85        """Save a fact to memory under a key, for recall later.
 86
 87        If the key already holds a value it is overwritten, so use a stable,
 88        descriptive key (e.g. "user_timezone", "project_deadline") and re-call
 89        remember() to update a fact.
 90
 91        Args:
 92            key: Short, stable identifier for the fact (used to recall it).
 93            value: The fact to store, as a self-contained string.
 94
 95        Returns:
 96            A confirmation that the fact was stored.
 97        """
 98        await self.store.set(key, value)
 99        self._notify(key, value)
100        return f"Remembered '{key}': {value}"
101
102    async def recall(self, key: str) -> str:
103        """Recall the value previously stored under a key.
104
105        Args:
106            key: The exact key the fact was stored under.
107
108        Returns:
109            The stored value, or a message saying nothing is stored under that
110            key. Use list_memories() if you are unsure of the exact key.
111        """
112        value = await self.store.get(key)
113        if value is None:
114            return f"No memory found under key '{key}'."
115        return value
116
117    async def forget(self, key: str) -> str:
118        """Delete the fact stored under a key.
119
120        Args:
121            key: The key to remove. Forgetting a key that does not exist is a
122                no-op and is reported as such.
123
124        Returns:
125            A confirmation describing what happened.
126        """
127        existed = await self.store.get(key) is not None
128        await self.store.delete(key)
129        self._notify(key, None)
130        if existed:
131            return f"Forgot '{key}'."
132        return f"Nothing was stored under '{key}'; nothing to forget."
133
134    async def list_memories(self) -> str:
135        """List every key and value currently held in memory.
136
137        Returns:
138            A newline-separated "key: value" listing, or a message saying
139            memory is empty.
140        """
141        items = await self.store.items()
142        if not items:
143            return "No memories stored."
144        return "\n".join(f"- {key}: {value}" for key, value in items.items())
145
146    async def list_keys(self) -> str:
147        """List just the keys currently held in memory, without their values.
148
149        Prefer this over list_memories() to see what is stored without pulling
150        every value into context; then recall(key) only the ones you need.
151
152        Returns:
153            A newline-separated list of keys, or a message saying memory is
154            empty.
155        """
156        keys = await self.store.keys()
157        if not keys:
158            return "No memories stored."
159        return "\n".join(f"- {key}" for key in keys)
160
161    async def search_memories(self, query: str) -> str:
162        """Search stored memories for relevance to a query across keys and values.
163
164        Use this when you remember roughly what a fact was about but not the
165        exact key. Ranking favors exact and substring key matches, but also
166        catches value hits, multi-word queries, and near-miss typos — so it
167        finds more than a plain substring search would.
168
169        Args:
170            query: Free-text search string.
171
172        Returns:
173            Matching "key: value" entries ranked by relevance, or a message
174            saying nothing matched.
175        """
176        items = await self.store.items()
177        hits = await self.search.search(items, query)
178        if not hits:
179            return f"No memories matched '{query}'."
180        return "\n".join(f"- {key}: {value}" for key, value, _score in hits)
181
182    @classmethod
183    def prompt(cls) -> str:
184        return (
185            "Use the memory tools to remember facts across the conversation. "
186            "Call remember(key, value) to save a fact under a short, stable, descriptive key; "
187            "re-calling remember() with the same key overwrites the old value. "
188            "Call recall(key) to read a fact back, and forget(key) to delete one. "
189            "If you are unsure of the exact key, call search_memories() to find related entries, "
190            "list_keys() to see what is stored without pulling in every value, "
191            "or list_memories() to see everything stored. "
192            "Save anything the user tells you that may be useful later (preferences, names, goals, "
193            "constraints), and recall before asking the user to repeat themselves."
194        )
195
196    def tool_set(self) -> list[RTFunction]:
197        functions = [
198            self.remember,
199            self.recall,
200            self.forget,
201            self.list_keys,
202            self.list_memories,
203            self.search_memories,
204        ]
205        return [rt.function_node(func) for func in functions]

 49class ToDoToolSet(ToolSet):
 50    def __init__(self, callback: Callable[[str, str, State], None] | None = None):
 51        """Create an empty toolset with an optional post-add callback.
 52
 53        Args:
 54            callback: Invoked after each successful add() with (short_description, description, state).
 55                      Exceptions are logged and swallowed; the todo is always committed regardless.
 56        """
 57        self.todos: list[ToDo] = []
 58        self._lock = asyncio.Lock()
 59
 60        if callback is None:
 61
 62            def default_add_callback(
 63                short_description: str, description: str, state: State
 64            ):
 65                pass
 66
 67            callback = default_add_callback
 68
 69        self.add_callback = callback
 70
 71    async def add(
 72        self, short_description: str, description: str, state: State = State.NOT_STARTED
 73    ):
 74        """Add a new todo to this toolset instance.
 75
 76        Args:
 77            short_description: Brief, unique label for the todo (used as its identifier in listings).
 78            description: Full details of what needs to be done.
 79            state: Initial state of the todo. Defaults to NOT_STARTED.
 80
 81        Raises:
 82            ValueError: If a todo with the same short_description or description already exists.
 83        """
 84        async with self._lock:
 85            to_do = ToDo(
 86                id=len(self.todos) + 1,
 87                short_description=short_description,
 88                description=description,
 89                state=state,
 90            )
 91
 92            validity_check = self.check_if_valid(self.todos, to_do)
 93            if validity_check is not None:
 94                raise ValueError(validity_check)
 95
 96            self.todos.append(to_do)
 97
 98        # Callback fires after the todo is committed; kept outside the lock so
 99        # user-provided callbacks cannot cause a deadlock.
100        try:
101            self.add_callback(short_description, description, state)
102        except Exception as e:
103            logger.error(f"Error in callback for todo: {e}")
104
105    @classmethod
106    def check_if_valid(cls, todos: list[ToDo], todo_to_add: ToDo) -> str | None:
107        """Return an error message if todo_to_add conflicts with existing todos, else None.
108
109        Args:
110            todos: The current list of todos to validate against.
111            todo_to_add: The candidate todo whose short_description and description are checked.
112        """
113        if todo_to_add.short_description in [todo.short_description for todo in todos]:
114            return f"Todo with short description '{todo_to_add.short_description}' already exists. Please provide a unique short description."
115
116        if todo_to_add.description in [todo.description for todo in todos]:
117            return f"Todo with description '{todo_to_add.description}' already exists. Please provide a unique description."
118
119        return None
120
121    def _get_all_todos(self) -> list[ToDo]:
122        return self.todos
123
124    async def get_all_todos(self) -> list[str]:
125        """Return complete_print() strings for all active (non-NO_LONGER_PLANNED) todos."""
126        async with self._lock:
127            return [
128                todo.complete_print()
129                for todo in self._get_all_todos()
130                if todo.state != State.NO_LONGER_PLANNED
131            ]
132
133    async def get_completed_todos(self) -> list[str]:
134        """Return complete_print() strings for todos in COMPLETED state."""
135        async with self._lock:
136            return [
137                todo.complete_print()
138                for todo in self._get_all_todos()
139                if todo.state == State.COMPLETED
140            ]
141
142    async def get_not_started_todos(self) -> list[str]:
143        """Return complete_print() strings for todos in NOT_STARTED state."""
144        async with self._lock:
145            return [
146                todo.complete_print()
147                for todo in self._get_all_todos()
148                if todo.state == State.NOT_STARTED
149            ]
150
151    async def get_incomplete_todos(self) -> list[str]:
152        """Return complete_print() strings for todos in NOT_STARTED, IN_PROGRESS, or FAILED state."""
153        incomplete_states = {State.NOT_STARTED, State.IN_PROGRESS, State.FAILED}
154        async with self._lock:
155            return [
156                todo.complete_print()
157                for todo in self._get_all_todos()
158                if todo.state in incomplete_states
159            ]
160
161    async def get_failed_todos(self) -> list[str]:
162        """Return complete_print() strings for todos in FAILED state."""
163        async with self._lock:
164            return [
165                todo.complete_print()
166                for todo in self._get_all_todos()
167                if todo.state == State.FAILED
168            ]
169
170    async def _find_and_update(self, todo_id: int, new_state: State) -> str:
171        async with self._lock:
172            for todo in self._get_all_todos():
173                if todo.id == todo_id:
174                    todo.update_state(new_state)
175                    return todo.complete_print()
176        raise ValueError(f"Todo with identifier '{todo_id}' not found.")
177
178    async def complete_todo_by_id(self, todo_id: int):
179        """Mark a todo as COMPLETED; raises ValueError if not found.
180
181        Args:
182            todo_id: The integer id of the todo to complete.
183        """
184        return "Successfully completed todo:\n" + await self._find_and_update(
185            todo_id, State.COMPLETED
186        )
187
188    async def start_todo_by_id(self, todo_id: int):
189        """Mark a todo as IN_PROGRESS; raises ValueError if not found.
190
191        Args:
192            todo_id: The integer id of the todo to start.
193        """
194        return "Successfully started todo:\n" + await self._find_and_update(
195            todo_id, State.IN_PROGRESS
196        )
197
198    async def fail_todo_by_id(self, todo_id: int):
199        """Mark a todo as FAILED; raises ValueError if not found.
200
201        Args:
202            todo_id: The integer id of the todo to fail.
203        """
204        return "Successfully marked todo as failed:\n" + await self._find_and_update(
205            todo_id, State.FAILED
206        )
207
208    async def no_longer_plan_todo_by_id(self, todo_id: int):
209        """Mark a todo as NO_LONGER_PLANNED; raises ValueError if not found.
210
211        Args:
212            todo_id: The integer id of the todo to deprioritize.
213        """
214        return (
215            "Successfully marked todo as no longer planned:\n"
216            + await self._find_and_update(todo_id, State.NO_LONGER_PLANNED)
217        )
218
219    async def make_all_no_longer_planned(self):
220        """Mark all NOT_STARTED and IN_PROGRESS todos as NO_LONGER_PLANNED; leaves COMPLETED and FAILED unchanged."""
221        affected = 0
222        async with self._lock:
223            for todo in self._get_all_todos():
224                if todo.state in {State.NOT_STARTED, State.IN_PROGRESS}:
225                    todo.update_state(State.NO_LONGER_PLANNED)
226                    affected += 1
227        return f"Marked {affected} todo(s) as no longer planned."
228
229    async def update_todo_by_id(self, todo_id: int, new_state: State):
230        """Transition a todo to an arbitrary state; raises ValueError if not found.
231
232        Args:
233            todo_id: The integer id of the todo to update.
234            new_state: The State to transition the todo to.
235        """
236        return "Successfully updated todo:\n" + await self._find_and_update(
237            todo_id, new_state
238        )
239
240    async def pretty_dashboard(self) -> str:
241        """Return a human-readable dashboard of active todos, or 'No todos found.'"""
242        async with self._lock:
243            lines = [
244                t.simplified_print()
245                for t in self._get_all_todos()
246                if t.state != State.NO_LONGER_PLANNED
247            ]
248        if not lines:
249            return "No todos found."
250        return "To-Dos\n" + "\n".join(lines)
251
252    @classmethod
253    def prompt(cls) -> str:
254        """Return the system prompt instructing an LLM how to use this toolset."""
255        return (
256            "Use the todo tools to plan and track your work. "
257            "Begin by calling add() for every task before starting any of them. "
258            "Call start_todo_by_id() when you begin a task and complete_todo_by_id() when it is done. "
259            "If a task cannot be completed, call fail_todo_by_id() instead. "
260            "If a planned task is no longer relevant, call no_longer_plan_todo_by_id() to remove it from active views. "
261            "To abandon the entire current plan, call make_all_no_longer_planned() — this leaves completed and failed todos unchanged. "
262            "Use update_todo_by_id() if a task needs a state change outside of the helpers above. "
263            "Retrieve identifiers via get_all_todos() before calling any id-based method. "
264            "Each todo requires a unique short_description and description."
265        )
266
267    def tool_set(self) -> list[RTFunction]:
268        """Return the list of RTFunction nodes for all public tools in this toolset."""
269        functions = [
270            self.add,
271            self.complete_todo_by_id,
272            self.start_todo_by_id,
273            self.fail_todo_by_id,
274            self.no_longer_plan_todo_by_id,
275            self.make_all_no_longer_planned,
276            self.update_todo_by_id,
277            self.get_all_todos,
278            self.get_completed_todos,
279            self.get_not_started_todos,
280            self.get_incomplete_todos,
281            self.get_failed_todos,
282        ]
283
284        return [rt.function_node(func) for func in functions]