1  #!/usr/bin/env python3
  2  """
  3  Todo Tool Module - Planning & Task Management
  4  
  5  Provides an in-memory task list the agent uses to decompose complex tasks,
  6  track progress, and maintain focus across long conversations. The state
  7  lives on the AIAgent instance (one per session) and is re-injected into
  8  the conversation after context compression events.
  9  
 10  Design:
 11  - Single `todo` tool: provide `todos` param to write, omit to read
 12  - Every call returns the full current list
 13  - No system prompt mutation, no tool response modification
 14  - Behavioral guidance lives entirely in the tool schema description
 15  """
 16  
 17  import json
 18  from typing import Dict, Any, List, Optional
 19  
 20  
 21  # Valid status values for todo items
 22  VALID_STATUSES = {"pending", "in_progress", "completed", "cancelled"}
 23  
 24  
 25  class TodoStore:
 26      """
 27      In-memory todo list. One instance per AIAgent (one per session).
 28  
 29      Items are ordered -- list position is priority. Each item has:
 30        - id: unique string identifier (agent-chosen)
 31        - content: task description
 32        - status: pending | in_progress | completed | cancelled
 33      """
 34  
 35      def __init__(self):
 36          self._items: List[Dict[str, str]] = []
 37  
 38      def write(self, todos: List[Dict[str, Any]], merge: bool = False) -> List[Dict[str, str]]:
 39          """
 40          Write todos. Returns the full current list after writing.
 41  
 42          Args:
 43              todos: list of {id, content, status} dicts
 44              merge: if False, replace the entire list. If True, update
 45                     existing items by id and append new ones.
 46          """
 47          if not merge:
 48              # Replace mode: new list entirely
 49              self._items = [self._validate(t) for t in self._dedupe_by_id(todos)]
 50          else:
 51              # Merge mode: update existing items by id, append new ones
 52              existing = {item["id"]: item for item in self._items}
 53              for t in self._dedupe_by_id(todos):
 54                  item_id = str(t.get("id", "")).strip()
 55                  if not item_id:
 56                      continue  # Can't merge without an id
 57  
 58                  if item_id in existing:
 59                      # Update only the fields the LLM actually provided
 60                      if "content" in t and t["content"]:
 61                          existing[item_id]["content"] = str(t["content"]).strip()
 62                      if "status" in t and t["status"]:
 63                          status = str(t["status"]).strip().lower()
 64                          if status in VALID_STATUSES:
 65                              existing[item_id]["status"] = status
 66                  else:
 67                      # New item -- validate fully and append to end
 68                      validated = self._validate(t)
 69                      existing[validated["id"]] = validated
 70                      self._items.append(validated)
 71              # Rebuild _items preserving order for existing items
 72              seen = set()
 73              rebuilt = []
 74              for item in self._items:
 75                  current = existing.get(item["id"], item)
 76                  if current["id"] not in seen:
 77                      rebuilt.append(current)
 78                      seen.add(current["id"])
 79              self._items = rebuilt
 80          return self.read()
 81  
 82      def read(self) -> List[Dict[str, str]]:
 83          """Return a copy of the current list."""
 84          return [item.copy() for item in self._items]
 85  
 86      def has_items(self) -> bool:
 87          """Check if there are any items in the list."""
 88          return bool(self._items)
 89  
 90      def format_for_injection(self) -> Optional[str]:
 91          """
 92          Render the todo list for post-compression injection.
 93  
 94          Returns a human-readable string to append to the compressed
 95          message history, or None if the list is empty.
 96          """
 97          if not self._items:
 98              return None
 99  
100          # Status markers for compact display
101          markers = {
102              "completed": "[x]",
103              "in_progress": "[>]",
104              "pending": "[ ]",
105              "cancelled": "[~]",
106          }
107  
108          # Only inject pending/in_progress items — completed/cancelled ones
109          # cause the model to re-do finished work after compression.
110          active_items = [
111              item for item in self._items
112              if item["status"] in ("pending", "in_progress")
113          ]
114          if not active_items:
115              return None
116  
117          lines = ["[Your active task list was preserved across context compression]"]
118          for item in active_items:
119              marker = markers.get(item["status"], "[?]")
120              lines.append(f"- {marker} {item['id']}. {item['content']} ({item['status']})")
121  
122          return "\n".join(lines)
123  
124      @staticmethod
125      def _validate(item: Dict[str, Any]) -> Dict[str, str]:
126          """
127          Validate and normalize a todo item.
128  
129          Ensures required fields exist and status is valid.
130          Returns a clean dict with only {id, content, status}.
131          """
132          item_id = str(item.get("id", "")).strip()
133          if not item_id:
134              item_id = "?"
135  
136          content = str(item.get("content", "")).strip()
137          if not content:
138              content = "(no description)"
139  
140          status = str(item.get("status", "pending")).strip().lower()
141          if status not in VALID_STATUSES:
142              status = "pending"
143  
144          return {"id": item_id, "content": content, "status": status}
145  
146      @staticmethod
147      def _dedupe_by_id(todos: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
148          """Collapse duplicate ids, keeping the last occurrence in its position."""
149          last_index: Dict[str, int] = {}
150          for i, item in enumerate(todos):
151              item_id = str(item.get("id", "")).strip() or "?"
152              last_index[item_id] = i
153          return [todos[i] for i in sorted(last_index.values())]
154  
155  
156  def todo_tool(
157      todos: Optional[List[Dict[str, Any]]] = None,
158      merge: bool = False,
159      store: Optional[TodoStore] = None,
160  ) -> str:
161      """
162      Single entry point for the todo tool. Reads or writes depending on params.
163  
164      Args:
165          todos: if provided, write these items. If None, read current list.
166          merge: if True, update by id. If False (default), replace entire list.
167          store: the TodoStore instance from the AIAgent.
168  
169      Returns:
170          JSON string with the full current list and summary metadata.
171      """
172      if store is None:
173          return tool_error("TodoStore not initialized")
174  
175      if todos is not None:
176          items = store.write(todos, merge)
177      else:
178          items = store.read()
179  
180      # Build summary counts
181      pending = sum(1 for i in items if i["status"] == "pending")
182      in_progress = sum(1 for i in items if i["status"] == "in_progress")
183      completed = sum(1 for i in items if i["status"] == "completed")
184      cancelled = sum(1 for i in items if i["status"] == "cancelled")
185  
186      return json.dumps({
187          "todos": items,
188          "summary": {
189              "total": len(items),
190              "pending": pending,
191              "in_progress": in_progress,
192              "completed": completed,
193              "cancelled": cancelled,
194          },
195      }, ensure_ascii=False)
196  
197  
198  def check_todo_requirements() -> bool:
199      """Todo tool has no external requirements -- always available."""
200      return True
201  
202  
203  # =============================================================================
204  # OpenAI Function-Calling Schema
205  # =============================================================================
206  # Behavioral guidance is baked into the description so it's part of the
207  # static tool schema (cached, never changes mid-conversation).
208  
209  TODO_SCHEMA = {
210      "name": "todo",
211      "description": (
212          "Manage your task list for the current session. Use for complex tasks "
213          "with 3+ steps or when the user provides multiple tasks. "
214          "Call with no parameters to read the current list.\n\n"
215          "Writing:\n"
216          "- Provide 'todos' array to create/update items\n"
217          "- merge=false (default): replace the entire list with a fresh plan\n"
218          "- merge=true: update existing items by id, add any new ones\n\n"
219          "Each item: {id: string, content: string, "
220          "status: pending|in_progress|completed|cancelled}\n"
221          "List order is priority. Only ONE item in_progress at a time.\n"
222          "Mark items completed immediately when done. If something fails, "
223          "cancel it and add a revised item.\n\n"
224          "Always returns the full current list."
225      ),
226      "parameters": {
227          "type": "object",
228          "properties": {
229              "todos": {
230                  "type": "array",
231                  "description": "Task items to write. Omit to read current list.",
232                  "items": {
233                      "type": "object",
234                      "properties": {
235                          "id": {
236                              "type": "string",
237                              "description": "Unique item identifier"
238                          },
239                          "content": {
240                              "type": "string",
241                              "description": "Task description"
242                          },
243                          "status": {
244                              "type": "string",
245                              "enum": ["pending", "in_progress", "completed", "cancelled"],
246                              "description": "Current status"
247                          }
248                      },
249                      "required": ["id", "content", "status"]
250                  }
251              },
252              "merge": {
253                  "type": "boolean",
254                  "description": (
255                      "true: update existing items by id, add new ones. "
256                      "false (default): replace the entire list."
257                  ),
258                  "default": False
259              }
260          },
261          "required": []
262      }
263  }
264  
265  
266  # --- Registry ---
267  from tools.registry import registry, tool_error
268  
269  registry.register(
270      name="todo",
271      toolset="todo",
272      schema=TODO_SCHEMA,
273      handler=lambda args, **kw: todo_tool(
274          todos=args.get("todos"), merge=args.get("merge", False), store=kw.get("store")),
275      check_fn=check_todo_requirements,
276      emoji="📋",
277  )