diff --git a/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py b/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py index 5df269e8c..7d8a24157 100644 --- a/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py +++ b/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py @@ -37,6 +37,11 @@ ) from idp_common.bedrock.client import CACHEPOINT_SUPPORTED_MODELS +from idp_common.utils.strands_agent_tools.todo_list import ( + create_todo_list, + update_todo, + view_todo_list, +) # Use AWS Lambda Powertools Logger for structured logging # Automatically logs as JSON with Lambda context, request_id, timestamp, etc. @@ -404,77 +409,6 @@ def patch_buffer_data(patches: list[dict[str, Any]], agent: Agent) -> str: return f"Successfully patched {str(patched_data)[100:]}...." -@tool -def create_todo_list(todos: list[str], agent: Agent) -> str: - """Create a new todo list to track your extraction tasks. Use this to plan your work, especially for large documents. - - Args: - todos: List of task descriptions to track (e.g., ["Extract rows 1-100", "Extract rows 101-200"]) - - Example: - create_todo_list(["Extract first 100 rows", "Extract rows 101-200", "Extract rows 201-300", "Validate and finalize"], agent) - """ - todo_list = [{"task": task, "completed": False} for task in todos] - agent.state.set("todo_list", todo_list) - logger.info("Created todo list", extra={"todo_count": len(todo_list)}) - return f"Created todo list with {len(todo_list)} tasks:\n" + "\n".join( - f"{i + 1}. [ ] {item['task']}" for i, item in enumerate(todo_list) - ) - - -@tool -def update_todo(task_index: int, completed: bool, agent: Agent) -> str: - """Mark a todo item as completed or not completed. - - Args: - task_index: Index of the task to update (1-based, matching the list display) - completed: True to mark as completed, False to mark as incomplete - - Example: - update_todo(1, True, agent) # Mark first task as completed - """ - todo_list: list[dict[str, Any]] | None = agent.state.get("todo_list") - - if not todo_list: - return "No todo list found. Create one first using create_todo_list." - - # Convert to 0-based index - index = task_index - 1 - - if index < 0 or index >= len(todo_list): - return f"Invalid task index {task_index}. Valid range: 1-{len(todo_list)}" - - todo_list[index]["completed"] = completed - agent.state.set("todo_list", todo_list) - - status = "completed" if completed else "incomplete" - logger.info( - f"Updated todo {task_index}", - extra={"task": todo_list[index]["task"], "completed": completed}, - ) - return f"Task {task_index} marked as {status}: {todo_list[index]['task']}" - - -@tool -def view_todo_list(agent: Agent) -> str: - """View your current todo list with completion status.""" - todo_list: list[dict[str, Any]] | None = agent.state.get("todo_list") - - if not todo_list: - return "No todo list found. Create one using create_todo_list to track your extraction tasks." - - completed_count = sum(1 for item in todo_list if item["completed"]) - total_count = len(todo_list) - - result = f"Todo List ({completed_count}/{total_count} completed):\n" - result += "\n".join( - f"{i + 1}. [{'✓' if item['completed'] else ' '}] {item['task']}" - for i, item in enumerate(todo_list) - ) - - return result - - SYSTEM_PROMPT = """ You are a useful assistant that helps turn unstructured data into structured data using the provided tools. diff --git a/lib/idp_common_pkg/idp_common/utils/strands_agent_tools/todo_list.py b/lib/idp_common_pkg/idp_common/utils/strands_agent_tools/todo_list.py new file mode 100644 index 000000000..1cf95aa8a --- /dev/null +++ b/lib/idp_common_pkg/idp_common/utils/strands_agent_tools/todo_list.py @@ -0,0 +1,79 @@ +import os +from aws_lambda_powertools import Logger +from typing_extensions import Any +from strands import Agent, tool + +logger = Logger(service="agentic_idp", level=os.getenv("LOG_LEVEL", "INFO")) + +@tool +def update_todo(task_index: int, completed: bool, agent: Agent) -> str: + """Mark a todo item as completed or not completed. + + Args: + task_index: Index of the task to update (1-based, matching the list display) + completed: True to mark as completed, False to mark as incomplete + + Example: + update_todo(1, True, agent) # Mark first task as completed + """ + todo_list: list[dict[str, Any]] | None = agent.state.get("todo_list") + + if not todo_list: + return "No todo list found. Create one first using create_todo_list." + + # Convert to 0-based index + index = task_index - 1 + + if index < 0 or index >= len(todo_list): + return f"Invalid task index {task_index}. Valid range: 1-{len(todo_list)}" + + todo_list[index]["completed"] = completed + agent.state.set("todo_list", todo_list) + + status = "completed" if completed else "incomplete" + logger.info( + f"Updated todo {task_index}", + extra={"task": todo_list[index]["task"], "completed": completed}, + ) + return f"Task {task_index} marked as {status}: {todo_list[index]['task']}" + + + + +@tool +def view_todo_list(agent: Agent) -> str: + """View your current todo list with completion status.""" + todo_list: list[dict[str, Any]] | None = agent.state.get("todo_list") + + if not todo_list: + return "No todo list found. Create one using create_todo_list to track your extraction tasks." + + completed_count = sum(1 for item in todo_list if item["completed"]) + total_count = len(todo_list) + + result = f"Todo List ({completed_count}/{total_count} completed):\n" + result += "\n".join( + f"{i + 1}. [{'✓' if item['completed'] else ' '}] {item['task']}" + for i, item in enumerate(todo_list) + ) + + return result + + + +@tool +def create_todo_list(todos: list[str], agent: Agent) -> str: + """Create a new todo list to track your extraction tasks. Use this to plan your work, especially for large documents. + + Args: + todos: List of task descriptions to track (e.g., ["Extract rows 1-100", "Extract rows 101-200"]) + + Example: + create_todo_list(["Extract first 100 rows", "Extract rows 101-200", "Extract rows 201-300", "Validate and finalize"], agent) + """ + todo_list = [{"task": task, "completed": False} for task in todos] + agent.state.set("todo_list", todo_list) + logger.info("Created todo list", extra={"todo_count": len(todo_list)}) + return f"Created todo list with {len(todo_list)} tasks:\n" + "\n".join( + f"{i + 1}. [ ] {item['task']}" for i, item in enumerate(todo_list) + )