railtracks.prebuilt

 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]

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

 31class WebSearchToolSet(ToolSet):
 32    """Prebuilt web search + page-fetch tools for an agent.
 33
 34    Gives an agent the ability to search the live web and read full page
 35    content
 36
 37    Args:
 38        search: Backend used by search()/search_and_fetch(). Defaults to
 39            ``TavilySearch()`` (requires a ``TAVILY_API_KEY``).
 40        fetch: Backend used by fetch()/search_and_fetch(). Defaults to
 41            ``HttpFetch()``.
 42    """
 43
 44    def __init__(
 45        self,
 46        search: SearchBackend | None = None,
 47        fetch: FetchBackend | None = None,
 48    ) -> None:
 49        self.search_backend: SearchBackend = (
 50            search if search is not None else _default_search()
 51        )
 52        self.fetch_backend: FetchBackend = (
 53            fetch if fetch is not None else _default_fetch()
 54        )
 55
 56    async def search(self, query: str, top_k: int = 5) -> str:
 57        """Search the web and return ranked results (title, URL, snippet).
 58
 59        Use this to find candidate pages before fetching full content with
 60        fetch(), or use search_and_fetch() to do both in one call.
 61
 62        Args:
 63            query: Free-text search query.
 64            top_k: Maximum number of results to return.
 65
 66        Returns:
 67            Newline-separated "title — url" entries with a snippet on each,
 68            ranked by relevance, or a message saying the search failed or
 69            found nothing.
 70        """
 71        try:
 72            results = await self.search_backend.search(query, top_k=top_k)
 73        except Exception as e:
 74            logger.error(f"WebSearch search backend error for query {query!r}: {e}")
 75            return f"Search failed: {e}"
 76
 77        if not results:
 78            return f"No results found for '{query}'."
 79        return "\n".join(f"- {r.title} — {r.url}\n  {r.snippet}" for r in results)
 80
 81    async def fetch(self, url: str) -> str:
 82        """Fetch a URL (typically from search() results) and return cleaned page text.
 83
 84        Args:
 85            url: The page URL to retrieve. Should be an http(s) URL, usually
 86                one returned by search().
 87
 88        Returns:
 89            The extracted, human-readable text content of the page, or a
 90            message describing why the fetch failed (blocked, paywalled,
 91            not found, no extractable content) so a different result can be
 92            tried instead.
 93        """
 94        try:
 95            result = await self.fetch_backend.fetch(url)
 96        except Exception as e:
 97            logger.error(f"WebSearch fetch backend error for url {url!r}: {e}")
 98            return f"Fetch failed for '{url}': {e}"
 99
100        if result.is_error:
101            return f"Fetch failed for '{url}': {result.error_message}"
102
103        header = f"{result.title}\n" if result.title else ""
104        return f"{header}{result.text}"
105
106    async def search_and_fetch(self, query: str, top_k: int = 3) -> str:
107        """Search the web and fetch full content for each top result in one call.
108
109        Convenience wrapper combining search() and fetch(); use when you want
110        full page content immediately without an extra round trip, at the
111        cost of fetching (and paying token cost for) top_k pages instead of
112        one. Individual fetch failures are reported inline rather than
113        aborting the whole call.
114
115        Args:
116            query: Free-text search query.
117            top_k: Number of top results to fetch full content for.
118
119        Returns:
120            For each result: title, url, and either its cleaned text or an
121            inline error message, separated by section dividers.
122        """
123        try:
124            results = await self.search_backend.search(query, top_k=top_k)
125        except Exception as e:
126            logger.error(f"WebSearch search backend error for query {query!r}: {e}")
127            return f"Search failed: {e}"
128
129        if not results:
130            return f"No results found for '{query}'."
131
132        sections = []
133        for r in results:
134            try:
135                fetched = await self.fetch_backend.fetch(r.url)
136                body = (
137                    fetched.text
138                    if not fetched.is_error
139                    else f"[fetch failed: {fetched.error_message}]"
140                )
141            except Exception as e:
142                logger.error(f"WebSearch fetch backend error for url {r.url!r}: {e}")
143                body = f"[fetch failed: {e}]"
144            sections.append(f"## {r.title}\n{r.url}\n\n{body}")
145        return "\n\n---\n\n".join(sections)
146
147    @classmethod
148    def prompt(cls) -> str:
149        return (
150            "Use the web search tools to find and read current information from the "
151            "live web. Call search(query) to get ranked title/url/snippet results, "
152            "then fetch(url) on the most promising result(s) to read full page content. "
153            "Use search_and_fetch(query) when you want full content for the top results "
154            "in a single step. If a fetch fails (blocked, paywalled, no content), try "
155            "another result instead of retrying the same URL."
156        )
157
158    def tool_set(self) -> list[RTFunction]:
159        functions = [self.search, self.fetch, self.search_and_fetch]
160        return [rt.function_node(func) for func in functions]