aip-agents-binary 0.5.20__py3-none-manylinux_2_31_x86_64.whl

aip_agents/memory/simple_memory.pyi

@@ -0,0 +1,23 @@
+from aip_agents.memory.base import BaseMemory as BaseMemory, ChatMessage as ChatMessage
+
+class SimpleMemory(BaseMemory):
+    """A simple memory implementation that stores messages in a list."""
+    messages: list[ChatMessage]
+    def __init__(self) -> None:
+        """Initialize the SimpleMemory instance with an empty message list."""
+    def add_message(self, message_or_role, content=None) -> None:
+        '''Add a message to memory.
+
+        Supports two calling patterns for backward compatibility:
+        1. add_message(ChatMessage) - Adds a ChatMessage object directly
+        2. add_message(role, content) - Creates and adds a ChatMessage with the given role and content
+
+        Args:
+            message_or_role: Either a ChatMessage object or a string role (e.g., "user", "assistant").
+            content: Optional content string when using the role+content pattern.
+                Required when message_or_role is a string role.
+        '''
+    def get_messages(self) -> list[ChatMessage]:
+        """Get all messages from memory."""
+    def clear(self) -> None:
+        """Clear all messages from memory."""

aip_agents/middleware/__init__.py

@@ -0,0 +1,17 @@
+"""Planning Middleware package.
+
+Provides composable middleware components for enhancing agent capabilities with
+planning and custom lifecycle hooks.
+"""
+
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+from aip_agents.middleware.manager import MiddlewareManager
+from aip_agents.middleware.todolist import TodoListMiddleware, TodoStatus
+
+__all__ = [
+    "AgentMiddleware",
+    "ModelRequest",
+    "MiddlewareManager",
+    "TodoListMiddleware",
+    "TodoStatus",
+]

aip_agents/middleware/__init__.pyi

@@ -0,0 +1,5 @@
+from aip_agents.middleware.base import AgentMiddleware as AgentMiddleware, ModelRequest as ModelRequest
+from aip_agents.middleware.manager import MiddlewareManager as MiddlewareManager
+from aip_agents.middleware.todolist import TodoListMiddleware as TodoListMiddleware, TodoStatus as TodoStatus
+
+__all__ = ['AgentMiddleware', 'ModelRequest', 'MiddlewareManager', 'TodoListMiddleware', 'TodoStatus']

aip_agents/middleware/base.py

@@ -0,0 +1,88 @@
+"""Base protocol and types for agent middleware.
+
+This module defines the core AgentMiddleware protocol that all middleware components
+must implement, along with the ModelRequest type for model invocations.
+"""
+
+from typing import Any, Protocol, TypedDict
+
+from langchain_core.tools import BaseTool
+
+
+class ModelRequest(TypedDict, total=False):
+    """Represents parameters for a model invocation that middleware can modify.
+
+    This TypedDict defines the structure of requests passed to the LLM, allowing
+    middleware to add tools or modify prompts before each invocation.
+
+    Attributes:
+        messages: List of messages in the conversation.
+        tools: List of tools available to the model.
+        system_prompt: System-level instruction for the model.
+    """
+
+    messages: list[Any]
+    tools: list[BaseTool]
+    system_prompt: str
+
+
+class AgentMiddleware(Protocol):
+    """Protocol defining the interface for composable agent middleware.
+
+    Middleware components can contribute tools, enhance system prompts, and provide
+    lifecycle hooks that execute before, during, and after model invocations.
+
+    All middleware must implement this protocol to be compatible with MiddlewareManager.
+
+    Attributes:
+        tools: List of tools contributed by this middleware.
+        system_prompt_additions: Optional text to append to the agent's system prompt.
+    """
+
+    tools: list[BaseTool]
+    system_prompt_additions: str | None
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed before each model invocation.
+
+        Use this hook to prepare state, log context, or perform setup tasks
+        before the model is called.
+
+        Args:
+            state: Current agent state containing messages and other context.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """
+        return {}  # pragma: no cover  # Protocol default - cannot be executed directly
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Hook to modify the model request before invocation.
+
+        Use this hook to add tools, modify the system prompt, adjust model parameters,
+        or change tool selection strategy.
+
+        Args:
+            request: The model request that will be sent to the LLM.
+            state: Current agent state for context.
+
+        Returns:
+            Modified ModelRequest. Can return the same request if no changes needed.
+        """
+        return request  # pragma: no cover  # Protocol default - cannot be executed directly  # pragma: no cover  # Protocol default - cannot be executed directly
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed after each model invocation.
+
+        Use this hook for cleanup, logging, state updates, or post-processing
+        of model outputs.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """
+        return {}  # pragma: no cover  # Protocol default - cannot be executed directly

aip_agents/middleware/base.pyi

@@ -0,0 +1,71 @@
+from langchain_core.tools import BaseTool as BaseTool
+from typing import Any, Protocol, TypedDict
+
+class ModelRequest(TypedDict, total=False):
+    """Represents parameters for a model invocation that middleware can modify.
+
+    This TypedDict defines the structure of requests passed to the LLM, allowing
+    middleware to add tools or modify prompts before each invocation.
+
+    Attributes:
+        messages: List of messages in the conversation.
+        tools: List of tools available to the model.
+        system_prompt: System-level instruction for the model.
+    """
+    messages: list[Any]
+    tools: list[BaseTool]
+    system_prompt: str
+
+class AgentMiddleware(Protocol):
+    """Protocol defining the interface for composable agent middleware.
+
+    Middleware components can contribute tools, enhance system prompts, and provide
+    lifecycle hooks that execute before, during, and after model invocations.
+
+    All middleware must implement this protocol to be compatible with MiddlewareManager.
+
+    Attributes:
+        tools: List of tools contributed by this middleware.
+        system_prompt_additions: Optional text to append to the agent's system prompt.
+    """
+    tools: list[BaseTool]
+    system_prompt_additions: str | None
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed before each model invocation.
+
+        Use this hook to prepare state, log context, or perform setup tasks
+        before the model is called.
+
+        Args:
+            state: Current agent state containing messages and other context.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Hook to modify the model request before invocation.
+
+        Use this hook to add tools, modify the system prompt, adjust model parameters,
+        or change tool selection strategy.
+
+        Args:
+            request: The model request that will be sent to the LLM.
+            state: Current agent state for context.
+
+        Returns:
+            Modified ModelRequest. Can return the same request if no changes needed.
+        """
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed after each model invocation.
+
+        Use this hook for cleanup, logging, state updates, or post-processing
+        of model outputs.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Dict of state updates to merge into the agent state. Return empty dict
+            if no updates are needed.
+        """

aip_agents/middleware/manager.py

@@ -0,0 +1,128 @@
+"""Middleware orchestration and lifecycle management.
+
+This module provides MiddlewareManager, which coordinates multiple middleware
+components and manages their lifecycle hooks.
+"""
+
+from typing import Any
+
+from langchain_core.tools import BaseTool
+
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+
+
+class MiddlewareManager:
+    """Orchestrates multiple middleware components and manages hook execution.
+
+    The manager collects tools from all middleware, builds enhanced system prompts,
+    and executes lifecycle hooks in the correct order (forward for setup, reverse
+    for cleanup).
+
+    Attributes:
+        middleware: List of middleware components in registration order.
+    """
+
+    def __init__(self, middleware: list[AgentMiddleware]) -> None:
+        """Initialize the middleware manager.
+
+        Args:
+            middleware: List of middleware components to manage. Order matters:
+                       hooks execute forward (first to last) for before/modify,
+                       and reverse (last to first) for after.
+        """
+        self.middleware = middleware
+
+    def get_all_tools(self) -> list[BaseTool]:
+        """Collect tools from all registered middleware.
+
+        Returns:
+            Combined list of all tools contributed by all middleware components.
+            Empty list if no middleware or if middleware provide no tools.
+        """
+        tools: list[BaseTool] = []
+        for mw in self.middleware:
+            tools.extend(mw.tools)
+        return tools
+
+    def build_system_prompt(self, base_instruction: str) -> str:
+        """Build enhanced system prompt by concatenating base instruction with middleware additions.
+
+        Args:
+            base_instruction: The base system prompt for the agent.
+
+        Returns:
+            Enhanced system prompt with all middleware additions appended.
+            If no middleware provide additions, returns base_instruction unchanged.
+        """
+        parts = [base_instruction]
+
+        for mw in self.middleware:
+            if mw.system_prompt_additions:
+                parts.append(mw.system_prompt_additions)
+
+        return "\n\n".join(parts)
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute before_model hooks for all middleware in forward order.
+
+        Hooks run first to last, allowing earlier middleware to prepare state
+        for later middleware.
+
+        Args:
+            state: Current agent state.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in order of execution.
+        """
+        state_updates: dict[str, Any] = {}
+
+        for mw in self.middleware:
+            mw_updates = mw.before_model(state)
+            if mw_updates:
+                state_updates.update(mw_updates)
+
+        return state_updates
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Execute modify_model_request hooks for all middleware in forward order.
+
+        Each middleware receives the request modified by previous middleware,
+        allowing them to build on each other's changes.
+
+        Args:
+            request: The model request to be modified.
+            state: Current agent state for context.
+
+        Returns:
+            Final modified request after all middleware have processed it.
+        """
+        current_request = request
+
+        for mw in self.middleware:
+            current_request = mw.modify_model_request(current_request, state)
+
+        return current_request
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute after_model hooks for all middleware in reverse order.
+
+        Hooks run last to first (reverse of registration order), allowing
+        proper cleanup and unwinding of middleware operations.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in reverse order of execution.
+        """
+        updates: dict[str, Any] = {}
+
+        # Execute in reverse order
+        for mw in reversed(self.middleware):
+            mw_updates = mw.after_model(state)
+            if mw_updates:
+                updates.update(mw_updates)
+
+        return updates

aip_agents/middleware/manager.pyi

@@ -0,0 +1,80 @@
+from _typeshed import Incomplete
+from aip_agents.middleware.base import AgentMiddleware as AgentMiddleware, ModelRequest as ModelRequest
+from langchain_core.tools import BaseTool as BaseTool
+from typing import Any
+
+class MiddlewareManager:
+    """Orchestrates multiple middleware components and manages hook execution.
+
+    The manager collects tools from all middleware, builds enhanced system prompts,
+    and executes lifecycle hooks in the correct order (forward for setup, reverse
+    for cleanup).
+
+    Attributes:
+        middleware: List of middleware components in registration order.
+    """
+    middleware: Incomplete
+    def __init__(self, middleware: list[AgentMiddleware]) -> None:
+        """Initialize the middleware manager.
+
+        Args:
+            middleware: List of middleware components to manage. Order matters:
+                       hooks execute forward (first to last) for before/modify,
+                       and reverse (last to first) for after.
+        """
+    def get_all_tools(self) -> list[BaseTool]:
+        """Collect tools from all registered middleware.
+
+        Returns:
+            Combined list of all tools contributed by all middleware components.
+            Empty list if no middleware or if middleware provide no tools.
+        """
+    def build_system_prompt(self, base_instruction: str) -> str:
+        """Build enhanced system prompt by concatenating base instruction with middleware additions.
+
+        Args:
+            base_instruction: The base system prompt for the agent.
+
+        Returns:
+            Enhanced system prompt with all middleware additions appended.
+            If no middleware provide additions, returns base_instruction unchanged.
+        """
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute before_model hooks for all middleware in forward order.
+
+        Hooks run first to last, allowing earlier middleware to prepare state
+        for later middleware.
+
+        Args:
+            state: Current agent state.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in order of execution.
+        """
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Execute modify_model_request hooks for all middleware in forward order.
+
+        Each middleware receives the request modified by previous middleware,
+        allowing them to build on each other's changes.
+
+        Args:
+            request: The model request to be modified.
+            state: Current agent state for context.
+
+        Returns:
+            Final modified request after all middleware have processed it.
+        """
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute after_model hooks for all middleware in reverse order.
+
+        Hooks run last to first (reverse of registration order), allowing
+        proper cleanup and unwinding of middleware operations.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in reverse order of execution.
+        """

aip_agents/middleware/todolist.py

@@ -0,0 +1,274 @@
+"""TodoList middleware for agent planning and task decomposition.
+
+Provides a write_todos tool that allows agents to break down complex tasks
+into discrete, trackable steps with status management.
+"""
+
+import threading
+from enum import StrEnum
+from typing import Any
+
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field, SkipValidation
+
+from aip_agents.middleware.base import ModelRequest
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class TodoStatus(StrEnum):
+    """Enumeration of possible todo item statuses."""
+
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+
+
+WRITE_TODOS_SYSTEM_PROMPT = """## Task Planning with `write_todos`
+
+Use `write_todos` to plan and track complex multi-step objectives (3+ steps, multiple deliverables, or tasks described as "comprehensive"). Skip for trivial single-action requests.
+
+**Execution rules:**
+1. Create the plan with `write_todos` BEFORE executing work
+2. Execute ALL tasks autonomously—do NOT pause to ask "should I continue?" or "would you like a preview?"
+3. Revise the plan as needed when new information emerges
+4. Mark tasks completed as you finish them
+5. **CRITICAL: Before your final response, call `write_todos` one last time to mark ALL tasks as `completed`. Never respond to the user while any task is `pending` or `in_progress`.**
+"""
+
+EMPTY_TODO_REMINDER = (
+    "\n\n**Todo List Reminder:** Your todo list is currently empty. "
+    "If you are working on tasks that would benefit from a todo list, "
+    "use the write_todos tool to create one. This helps track progress "
+    "and ensures you don't forget important steps.\n"
+)
+
+
+class TodoItem(BaseModel):
+    """Represents a single todo item in the agent's plan.
+
+    Attributes:
+        content: Human-readable description of the task.
+        active_form: Short imperative phrase for UI display.
+        status: Current status of the todo item.
+    """
+
+    content: str = Field(description="Description of what needs to be done")
+    active_form: str = Field(description="Short active form for display (e.g., 'researching topic')")
+    status: TodoStatus = Field(default=TodoStatus.PENDING, description="Current status of this todo")
+
+
+class TodoList(BaseModel):
+    """Represents a complete todo list for a thread.
+
+    Attributes:
+        items: List of todo items in order.
+    """
+
+    items: list[TodoItem] = Field(default_factory=list)
+
+
+WRITE_TODOS_TOOL_DESCRIPTION = """Create and manage a structured task list for complex work sessions (3+ steps).
+
+## When to Use
+- Multi-step tasks with 3+ distinct steps
+- Requests with multiple deliverables
+- User provides a list of tasks
+- User explicitly requests a todo list
+
+## When to Skip
+- Single straightforward action
+- Trivial tasks with no tracking benefit
+
+## Task States
+- `pending`: Not started
+- `in_progress`: Currently working
+- `completed`: Fully finished
+
+## Workflow
+1. **Plan first**: Call `write_todos` before executing work. Mark first task(s) as `in_progress`.
+2. **Stay autonomous**: Execute all tasks without pausing for user confirmation.
+3. **Adapt the plan**: Add new tasks as discovered, remove irrelevant ones.
+4. **Mark progress**: Update task status as you complete work.
+5. **Final update**: Before responding to user, call `write_todos` to mark ALL remaining tasks `completed`.
+
+## CRITICAL: Final Response Rule
+You must NOT respond to the user until EVERY task is marked `completed`. Your final action before responding should be a `write_todos` call that marks the last task(s) complete. If any task remains `pending` or `in_progress`, you are not done—keep working or update the plan."""
+
+
+class WriteTodosInput(BaseModel):
+    """Input schema for the write_todos tool."""
+
+    todos: TodoList = Field(
+        ...,
+        description="List of todo items, each with 'content' and 'activeForm' keys.",
+    )
+
+
+class WriteTodosTool(BaseTool):
+    """LangChain-compatible tool for managing todo lists via TodoListMiddleware."""
+
+    name: str = "write_todos_tool"
+    description: str = WRITE_TODOS_TOOL_DESCRIPTION
+    args_schema: type[BaseModel] = WriteTodosInput
+
+    # Instance reference to middleware storage (excluded from serialization, skip validation)
+    storage: SkipValidation[dict[str, TodoList]] = Field(default_factory=dict, exclude=True)
+    storage_lock: SkipValidation[threading.RLock] = Field(default_factory=threading.RLock, exclude=True)
+
+    def _run(self, todos: TodoList, config: RunnableConfig | None = None) -> str:  # type: ignore[override]
+        """Synchronous entrypoint used by generic tool runners.
+
+        Args:
+            todos: The todo list to write.
+            config: Optional configuration for the tool run.
+
+        Returns:
+            str: A success message indicating the number of todo items written.
+        """
+        tid = "default"
+        if config and "configurable" in config:
+            tid = config["configurable"].get("thread_id", "default")
+
+        try:
+            # Store in thread-safe manner
+            with self.storage_lock:
+                self.storage[tid] = todos
+
+            return f"Successfully wrote {len(todos.items)} todo items to your plan."
+        except Exception as e:  # pragma: no cover - defensive
+            return f"Error writing todos: {str(e)}"
+
+
+class TodoListMiddleware:
+    """Middleware that provides planning capabilities via todo list management.
+
+    Adds the write_todos tool and enhances the system prompt with planning
+    instructions, encouraging agents to break down complex tasks.
+
+    This middleware maintains thread-isolated todo lists, ensuring that
+    different conversation threads don't interfere with each other.
+
+    Each middleware instance has its own storage, preventing race conditions
+    when multiple agent instances are used concurrently.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the TodoList middleware with planning tools and instructions."""
+        # Instance-level storage to prevent race conditions across agent instances
+        self._storage: dict[str, TodoList] = {}
+        self._storage_lock = threading.RLock()
+
+        # Create tool with reference to this instance's storage
+        self.tools = [WriteTodosTool(storage=self._storage, storage_lock=self._storage_lock)]
+        self.system_prompt_additions = WRITE_TODOS_SYSTEM_PROMPT
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed before model invocation.
+
+        Syncs todos FROM state TO internal storage for LangGraph agents.
+        This allows todos to be persisted via LangGraph checkpointer.
+
+        Args:
+            state: Current agent state (may contain 'todos' key).
+
+        Returns:
+            Empty dict (no state updates needed).
+        """
+        # Sync from LangGraph state to internal storage (if state has todos)
+        if "todos" in state and state["todos"] is not None:
+            thread_id = state.get("thread_id", "default")
+
+            # Deep copy to prevent mutation issues between state and storage
+            with self._storage_lock:
+                self._storage[thread_id] = state["todos"].model_copy(deep=True)
+
+        return {}
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Hook to modify model request before invocation.
+
+        Injects current todo list status into the system prompt, ensuring
+        the agent has visibility into its current plan on every turn.
+        This follows Claude Code's pattern of injecting system reminders.
+
+        Args:
+            request: The model request.
+            state: Current agent state.
+
+        Returns:
+            Modified request with todo status injected into system prompt.
+        """
+        thread_id = state.get("thread_id", "default")
+
+        # Get current todos from storage
+        with self._storage_lock:
+            current_todos = self._storage.get(thread_id)
+
+        # Inject todo status into system prompt
+        if current_todos and current_todos.items:
+            # Build todo status reminder
+            todo_reminder = "\n\n## Current Todo List Status\n\n"
+            todo_reminder += "Your current todo list:\n"
+
+            for item in current_todos.items:
+                todo_reminder += f"[{item.status.value.upper()}] {item.content}\n"
+
+            todo_reminder += "\n**Remember to update the todo list as you complete tasks using write_todos.**\n"
+
+            # Append to system prompt
+            current_system_prompt = request.get("system_prompt", "")
+            request["system_prompt"] = current_system_prompt + todo_reminder
+        else:
+            # Static reminder when no todos exist
+            current_system_prompt = request.get("system_prompt", "")
+            request["system_prompt"] = current_system_prompt + EMPTY_TODO_REMINDER
+
+        return request
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Hook executed after model invocation.
+
+        Syncs todos FROM internal storage TO state for LangGraph agents.
+        This ensures any tool updates are reflected in the checkpointed state.
+
+        Args:
+            state: Current agent state.
+
+        Returns:
+            Dict with 'todos' key containing updated TodoList, or empty dict.
+        """
+        thread_id = state.get("thread_id", "default")
+
+        with self._storage_lock:
+            if thread_id in self._storage:
+                # Deep copy to prevent mutation issues between storage and state
+                return {"todos": self._storage[thread_id].model_copy(deep=True)}
+
+        return {}
+
+    def get_todos(self, thread_id: str = "default") -> TodoList:
+        """Retrieve the todo list for a specific thread.
+
+        Args:
+            thread_id: Thread identifier. Defaults to "default".
+
+        Returns:
+            TodoList for the thread, or empty list if none exists.
+        """
+        with self._storage_lock:
+            return self._storage.get(thread_id, TodoList())
+
+    def clear_todos(self, thread_id: str = "default") -> None:
+        """Clear the todo list for a specific thread.
+
+        Useful for cleanup between test runs or conversation resets.
+
+        Args:
+            thread_id: Thread identifier. Defaults to "default".
+        """
+        with self._storage_lock:
+            if thread_id in self._storage:
+                del self._storage[thread_id]