aip-agents-binary 0.5.25b8__py3-none-any.whl → 0.6.0__py3-none-any.whl

aip_agents/mcp/client/transports.py

@@ -12,10 +12,11 @@ from collections.abc import AsyncIterator
 from enum import StrEnum
 from typing import Any, Protocol
 
+import httpx
 from gllm_tools.mcp.client.config import MCPConfiguration
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters, stdio_client
-from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.streamable_http import streamable_http_client
 
 from aip_agents.utils.logger import get_logger
 
@@ -142,6 +143,27 @@ class SSETransport(Transport):
 class HTTPTransport(Transport):
     """Streamable HTTP transport handler."""
 
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize the HTTP transport.
+
+        Args:
+            server_name (str): Name of the MCP server.
+            config (MCPConfiguration): Configuration for the transport.
+        """
+        super().__init__(server_name, config)
+        self._http_client: httpx.AsyncClient | None = None
+
+    async def close(self) -> None:
+        """Clean up the transport connection and any owned HTTP client."""
+        await super().close()
+        if self._http_client:
+            try:
+                await self._http_client.aclose()
+            except Exception as e:
+                logger.warning(f"Error during HTTP client cleanup for {self.server_name}: {e}")
+            finally:
+                self._http_client = None
+
     async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
         """Connect using streamable HTTP transport.
 
@@ -163,11 +185,20 @@ class HTTPTransport(Transport):
         headers = _sanitize_headers(self.config)
         logger.debug(f"Attempting streamable HTTP connection to {url} with headers: {list(headers.keys())}")
         try:
-            self.ctx = streamablehttp_client(url=url, timeout=timeout, headers=headers)
+            http_client = httpx.AsyncClient(timeout=httpx.Timeout(timeout), headers=headers)
+            self._http_client = http_client
+            self.ctx = streamable_http_client(url=url, http_client=http_client)
             read_stream, write_stream, _ = await self.ctx.__aenter__()
             logger.info(f"Connected to {self.server_name} via HTTP")
             return read_stream, write_stream, self.ctx
         except Exception as e:
+            if self._http_client:
+                try:
+                    await self._http_client.aclose()
+                except Exception as close_exc:
+                    logger.warning(f"Error during HTTP client cleanup for {self.server_name}: {close_exc}")
+                finally:
+                    self._http_client = None
             raise ConnectionError(f"HTTP connection failed for {self.server_name}: {str(e)}") from e
 
 

aip_agents/mcp/client/transports.pyi

@@ -0,0 +1,132 @@
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from aip_agents.utils.logger import get_logger as get_logger
+from collections.abc import AsyncIterator as AsyncIterator
+from enum import StrEnum
+from gllm_tools.mcp.client.config import MCPConfiguration
+from typing import Any, Protocol
+
+class TransportContext(Protocol):
+    """Protocol defining the interface for async context managers used in MCP transport connections."""
+    async def __aenter__(self) -> None:
+        """Enter the async context, establishing the connection and returning read/write streams."""
+    async def __aexit__(self, _exc_type, _exc_val, _exc_tb) -> None:
+        """Exit the async context, performing cleanup and closing the connection.
+
+        Args:
+            _exc_type: Exception type if an exception occurred.
+            _exc_val: Exception value if an exception occurred.
+            _exc_tb: Exception traceback if an exception occurred.
+        """
+
+logger: Incomplete
+DEFAULT_TIMEOUT: float
+
+class TransportType(StrEnum):
+    """Enum for supported MCP transport types."""
+    HTTP: str
+    SSE: str
+    STDIO: str
+
+class Transport(ABC):
+    """Abstract base class for MCP transports."""
+    server_name: Incomplete
+    config: Incomplete
+    ctx: Any
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize the transport.
+
+        Args:
+            server_name (str): Name of the MCP server.
+            config (MCPConfiguration): Configuration for the transport.
+        """
+    @abstractmethod
+    async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
+        """Establish connection and return read/write streams and context manager.
+
+        Returns:
+            tuple[AsyncIterator[bytes], AsyncIterator[bytes], Any]:
+                (read_stream, write_stream, ctx)
+            Where:
+                - read_stream: AsyncIterator[bytes] for reading from the server.
+                - write_stream: AsyncIterator[bytes] for writing to the server.
+                - ctx: The async context manager instance for cleanup via __aexit__.
+
+        Raises:
+            ValueError: If required config (e.g., URL or command) is missing.
+            ConnectionError: If connection establishment fails.
+        """
+    async def close(self) -> None:
+        """Clean up the transport connection."""
+
+class SSETransport(Transport):
+    """SSE transport handler."""
+    ctx: Incomplete
+    async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
+        """Connect using SSE transport.
+
+        Builds SSE URL from config, initializes client with timeout, and enters context.
+
+        Returns:
+            tuple[AsyncIterator[bytes], AsyncIterator[bytes], Any]: (read_stream, write_stream, ctx)
+
+        Raises:
+            ValueError: If URL is missing.
+            ConnectionError: If SSE connection fails.
+        """
+
+class HTTPTransport(Transport):
+    """Streamable HTTP transport handler."""
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize the HTTP transport.
+
+        Args:
+            server_name (str): Name of the MCP server.
+            config (MCPConfiguration): Configuration for the transport.
+        """
+    async def close(self) -> None:
+        """Clean up the transport connection and any owned HTTP client."""
+    ctx: Incomplete
+    async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
+        """Connect using streamable HTTP transport.
+
+        Builds MCP URL from config, initializes client with timeout, and enters context.
+
+        Returns:
+            tuple[AsyncIterator[bytes], AsyncIterator[bytes], Any]: (read_stream, write_stream, ctx)
+
+        Raises:
+            ValueError: If URL is missing.
+            ConnectionError: If HTTP connection fails.
+        """
+
+class StdioTransport(Transport):
+    """STDIO transport handler."""
+    ctx: Incomplete
+    async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
+        """Connect using STDIO transport.
+
+        Initializes stdio client from command/args/env in config and enters context.
+
+        Returns:
+            tuple[AsyncIterator[bytes], AsyncIterator[bytes], Any]: (read_stream, write_stream, ctx)
+
+        Raises:
+            ValueError: If command is missing.
+            ConnectionError: If STDIO connection fails.
+        """
+
+def create_transport(server_name: str, config: MCPConfiguration, transport_type: TransportType | str) -> Transport:
+    """Factory to create the appropriate transport instance.
+
+    Args:
+        server_name (str): Server name
+        config (MCPConfiguration): Config
+        transport_type (str): Transport type ('http', 'sse', 'stdio')
+
+    Returns:
+        Transport: Concrete transport instance
+
+    Raises:
+        ValueError: If transport_type is unsupported.
+    """

aip_agents/mcp/utils/config_validator.pyi

@@ -0,0 +1,82 @@
+from _typeshed import Incomplete
+from aip_agents.utils.logger import get_logger as get_logger
+from typing import Any
+
+logger: Incomplete
+
+def validate_allowed_tools_list(allowed_tools: list[str] | None, context: str) -> list[str] | None:
+    '''Validate that allowed_tools is a list of strings or None.
+
+    This function validates the type and contents of allowed_tools parameter.
+    It can be used to validate allowed_tools from any source (config dict, function parameter, etc).
+
+    Args:
+        allowed_tools: The allowed_tools value to validate (can be any type)
+        context: Context string for error messages (e.g., "Server \'my_server\'", "\'allowed_tools\' parameter")
+
+    Returns:
+        Validated list of allowed tool names, or None if allowed_tools is None or empty.
+        None/empty means no restriction (all tools allowed).
+
+    Raises:
+        ValueError: If allowed_tools is not None/list or contains non-string elements
+
+    Examples:
+        >>> validate_allowed_tools_list(None, "test")
+        None
+        >>> validate_allowed_tools_list([], "test")
+        None
+        >>> validate_allowed_tools_list(["tool1", "tool2"], "test")
+        [\'tool1\', \'tool2\']
+        >>> validate_allowed_tools_list("invalid", "test")  # doctest: +SKIP
+        ValueError: test: \'allowed_tools\' must be a list of strings, got str
+    '''
+def validate_allowed_tools_config(config: dict[str, Any], server_name: str) -> list[str] | None:
+    '''Validate and extract allowed_tools configuration from server config.
+
+    This function validates that the allowed_tools field, if present, is a list of strings.
+    It returns a normalized list of allowed tool names, or None if not specified/empty.
+
+    Args:
+        config: Server configuration dictionary that may contain \'allowed_tools\' field
+        server_name: Name of the server (for error messages)
+
+    Returns:
+        List of allowed tool names, or None if allowed_tools is not specified or is empty.
+        None means no restriction (all tools allowed).
+
+    Raises:
+        ValueError: If allowed_tools is present but not a list
+        ValueError: If allowed_tools contains non-string elements
+
+    Examples:
+        >>> validate_allowed_tools_config({"url": "..."}, "my_server")
+        None
+        >>> validate_allowed_tools_config({"url": "...", "allowed_tools": []}, "my_server")
+        None
+        >>> validate_allowed_tools_config({"url": "...", "allowed_tools": ["tool1", "tool2"]}, "my_server")
+        [\'tool1\', \'tool2\']
+    '''
+def validate_mcp_server_config(config: dict[str, Any], server_name: str) -> dict[str, Any]:
+    '''Validate complete MCP server configuration including allowed_tools.
+
+    This function performs comprehensive validation on an MCP server configuration,
+    ensuring all required fields are present and allowed_tools (if present) is valid.
+    Unknown configuration fields (such as any leftover disabled_tools) are silently ignored.
+
+    Args:
+        config: Server configuration dictionary
+        server_name: Name of the server (for error messages)
+
+    Returns:
+        Validated configuration dictionary (same as input, after validation)
+
+    Raises:
+        ValueError: If configuration is invalid
+
+    Examples:
+        >>> validate_mcp_server_config({"url": "http://localhost:8080"}, "my_server")
+        {\'url\': \'http://localhost:8080\'}
+        >>> validate_mcp_server_config({"command": "python", "args": ["server.py"]}, "my_server")
+        {\'command\': \'python\', \'args\': [\'server.py\']}
+    '''

aip_agents/memory/__init__.pyi

@@ -0,0 +1,5 @@
+from aip_agents.memory.base import BaseMemory as BaseMemory
+from aip_agents.memory.constants import MemoryMethod as MemoryMethod
+from aip_agents.memory.factory import MemoryFactory as MemoryFactory
+
+__all__ = ['BaseMemory', 'MemoryMethod', 'MemoryFactory']

aip_agents/memory/adapters/__init__.pyi

@@ -0,0 +1,4 @@
+from aip_agents.memory.adapters.base_adapter import BaseMemoryAdapter as BaseMemoryAdapter
+from aip_agents.memory.adapters.mem0 import Mem0Memory as Mem0Memory
+
+__all__ = ['Mem0Memory', 'BaseMemoryAdapter']

aip_agents/memory/adapters/base_adapter.pyi

@@ -0,0 +1,150 @@
+from _typeshed import Incomplete
+from aip_agents.memory.base import BaseMemory as BaseMemory
+from aip_agents.memory.constants import MemoryDefaults as MemoryDefaults
+from aip_agents.types import ChatMessage as ChatMessage
+from aip_agents.utils.datetime import format_created_updated_label as format_created_updated_label
+from aip_agents.utils.logger import get_logger as get_logger
+from concurrent.futures import Future
+from dataclasses import dataclass
+from gllm_memory import MemoryManager
+from gllm_memory.enums import MemoryScope
+from typing import Any, ClassVar
+
+MemoryManager = Any
+MemoryScope = Any
+logger: Incomplete
+DEFAULT_SCOPE: ClassVar[set[MemoryScope]]
+
+@dataclass(frozen=True)
+class _RetrieveOptions:
+    user_id: str
+    top_k: int
+    metadata: dict[str, str] | None
+    keywords: Any
+    page: int
+    categories: list[str] | None
+
+class _AsyncRunner:
+    """Runs async coroutines on a dedicated background event loop."""
+    def __init__(self) -> None: ...
+    def run(self, awaitable: Any) -> Any:
+        """Execute an awaitable on the background loop and block for the result.
+
+        Args:
+            awaitable: The coroutine to execute.
+
+        Returns:
+            The result of the coroutine execution.
+        """
+    def shutdown(self, timeout: float | None = 5.0) -> None:
+        """Gracefully stop the event loop thread.
+
+        Args:
+            timeout: Maximum time to wait for the thread to stop.
+        """
+    def __del__(self) -> None: ...
+
+class BaseMemoryAdapter(BaseMemory):
+    """Provider-agnostic long-term memory adapter backed by gllm_memory."""
+    agent_id: Incomplete
+    namespace: Incomplete
+    limit: Incomplete
+    max_chars: Incomplete
+    def __init__(self, *, agent_id: str, manager: MemoryManager, namespace: str | None = None, limit: int = ..., max_chars: int = ...) -> None:
+        """Initialize the GLLM memory adapter.
+
+        Args:
+            agent_id: Unique identifier for the agent using this memory.
+            manager: Configured gllm_memory MemoryManager instance.
+            namespace: Optional namespace for organizing memories.
+            limit: Maximum number of memories to retrieve in search operations.
+            max_chars: Maximum character length for text content.
+        """
+    @classmethod
+    def validate_env(cls) -> None:
+        """Base adapter does not enforce environment validation."""
+    def get_messages(self) -> list[ChatMessage]:
+        """Retrieve all stored chat messages.
+
+        Returns:
+            An empty list as GLLM adapter doesn't support message retrieval.
+        """
+    def add_message(self, message: ChatMessage) -> None:
+        """Best-effort single-message persistence for API parity.
+
+        Args:
+            message: The chat message to add to memory.
+        """
+    def clear(self) -> None:
+        """Clear all stored memories.
+
+        Raises:
+            NotImplementedError: This method is not implemented for GLLM adapter.
+        """
+    def search(self, query: str, *, user_id: str, limit: int | None = None, filters: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """Search for memories using a text query.
+
+        Args:
+            query: The search query string.
+            user_id: User identifier for the search scope.
+            limit: Maximum number of results to return.
+            filters: Optional filters to apply to the search.
+
+        Returns:
+            List of memory hits matching the search criteria.
+        """
+    def retrieve(self, *, query: str | None, user_id: str, limit: int | None = None, filters: dict[str, Any] | None = None, page: int | None = None) -> list[dict[str, Any]]:
+        """Retrieve memories with optional search query and filters.
+
+        Args:
+            query: Optional search query string. If None, retrieves all memories.
+            user_id: User identifier for the retrieval scope.
+            limit: Maximum number of results to return.
+            filters: Optional filters to apply to the retrieval.
+            page: Page number for pagination.
+
+        Returns:
+            List of memory hits matching the criteria.
+        """
+    def save_interaction(self, *, user_text: str, ai_text: str, user_id: str) -> None:
+        """Save a user-AI interaction as memories.
+
+        Args:
+            user_text: The user's input text.
+            ai_text: The AI's response text.
+            user_id: User identifier for the memory storage.
+        """
+    def save_interaction_async(self, *, user_text: str, ai_text: str, user_id: str) -> Future[Any]:
+        """Schedule save_interaction without blocking the caller.
+
+        Args:
+            user_text: The user's input text to save.
+            ai_text: The AI's response text to save.
+            user_id: User identifier for the memory storage.
+        """
+    def format_hits(self, hits: list[dict[str, Any]], max_items: int = ..., with_tag: bool = True) -> str:
+        """Format memory hits into a readable string.
+
+        Args:
+            hits: List of memory hit dictionaries to format.
+            max_items: Maximum number of hits to include in the output.
+            with_tag: Whether to wrap the output with memory tags.
+
+        Returns:
+            Formatted string representation of the memory hits.
+        """
+    def flush_pending_writes(self, timeout: float | None = None) -> None:
+        """Block until current async writes complete.
+
+        Args:
+            timeout: Maximum time to wait for pending writes to complete.
+        """
+    def close(self, *, wait: bool = True, timeout: float | None = None) -> None:
+        """Release background resources and optionally wait for pending saves.
+
+        Args:
+            wait: Whether to wait for pending async operations to complete.
+            timeout: Maximum time to wait when wait=True.
+        """
+    def __del__(self) -> None:
+        """Clean up resources when the adapter is garbage collected."""

aip_agents/memory/adapters/mem0.pyi

@@ -0,0 +1,22 @@
+from aip_agents.memory.adapters.base_adapter import BaseMemoryAdapter as BaseMemoryAdapter
+from aip_agents.memory.constants import MemoryDefaults as MemoryDefaults
+from typing import Any
+
+MemoryManager = Any
+
+class Mem0Memory(BaseMemoryAdapter):
+    """Mem0-backed long-term memory adapter using the gllm_memory SDK."""
+    def __init__(self, *, agent_id: str, namespace: str | None = None, limit: int = ..., max_chars: int = ..., host: str | None = None, instruction: str | None = None) -> None:
+        """Initialize the Mem0 memory adapter.
+
+        Args:
+            agent_id: Unique identifier for the agent using this memory.
+            namespace: Optional namespace for organizing memories.
+            limit: Maximum number of memories to retrieve in search operations.
+            max_chars: Maximum character length for text content.
+            host: Optional host URL for the Mem0 service.
+            instruction: Optional instruction text for memory operations.
+        """
+    @classmethod
+    def validate_env(cls) -> None:
+        """Ensure the Mem0 API key is available."""

aip_agents/memory/base.pyi

@@ -0,0 +1,60 @@
+from abc import ABC, abstractmethod
+from aip_agents.types import ChatMessage as ChatMessage
+from collections.abc import Sequence
+from typing import Any
+
+class BaseMemory(ABC):
+    """Base class for agent memory.
+
+    This concrete base class provides a default structure. Subclasses
+    can inherit from this class to implement specific memory management
+    behaviors.
+    """
+    @classmethod
+    def validate_env(cls) -> None:
+        """Validate environment prerequisites for a memory backend.
+
+        This hook allows memory implementations to fail fast when required
+        environment variables or credentials are missing. Default is a no-op.
+        """
+    @abstractmethod
+    def get_messages(self) -> list[ChatMessage]:
+        """Retrieve a list of messages.
+
+        Returns:
+            List[ChatMessage]: A list of messages in a generic format.
+        """
+    @abstractmethod
+    def add_message(self, message: ChatMessage) -> None:
+        """Add a message to the memory.
+
+        Adds a single ChatMessage to the memory storage. The exact implementation
+        depends on the specific memory backend.
+
+        Args:
+            message: The ChatMessage object to add to memory. The message should
+                contain role and content information.
+        """
+    def add_messages(self, messages: Sequence[ChatMessage]) -> None:
+        """Add multiple messages to the memory.
+
+        Args:
+            messages: A sequence of ChatMessage objects to add.
+        """
+    @abstractmethod
+    def clear(self) -> None:
+        """Clears the memory or resets the state of the agent.
+
+        This method must be implemented to define the specific behavior
+        for clearing or resetting the memory of the agent.
+        """
+    def get_memory_variables(self) -> dict[str, Any]:
+        """Retrieve memory variables.
+
+        This method returns a dictionary containing memory-related variables.
+        The default implementation returns a dictionary with chat_history.
+
+        Returns:
+            Dict[str, Any]: A dictionary where keys are variable names and values
+            are the corresponding memory-related data.
+        """

aip_agents/memory/constants.pyi

@@ -0,0 +1,25 @@
+from _typeshed import Incomplete
+
+class MemoryMethod:
+    """Constants for memory method names used in hasattr checks."""
+    SEARCH: str
+    SAVE_INTERACTION: str
+    FORMAT_HITS: str
+
+class MemoryDefaults:
+    """Default values for memory configuration parameters."""
+    MAX_ITEMS: int
+    RETRIEVAL_LIMIT: int
+    MAX_CHARS: int
+    LOG_PREVIEW_LENGTH: int
+    SMALL_PREVIEW_LENGTH: int
+    AGENT_ID_PREFIX: str
+    MEMORY_TAG_OPEN: str
+    MEMORY_TAG_CLOSE: str
+    DEFAULT_USER_ID: str
+    DATE_STRING_LENGTH: int
+
+class MemoryBackends:
+    """Supported memory backend identifiers."""
+    MEM0: str
+    SUPPORTED: Incomplete

aip_agents/memory/factory.pyi

@@ -0,0 +1,24 @@
+from _typeshed import Incomplete
+from aip_agents.memory.base import BaseMemory as BaseMemory
+from aip_agents.utils.logger import get_logger as get_logger
+from typing import Any
+
+logger: Incomplete
+BACKENDS: Incomplete
+
+class MemoryFactory:
+    """Factory to build concrete memory adapters by backend name."""
+    @staticmethod
+    def create(backend: str, **kwargs: Any) -> BaseMemory:
+        '''Create a memory adapter instance.
+
+        Args:
+            backend: Backend identifier (e.g., "mem0").
+            **kwargs: Keyword args passed to adapter constructor (e.g., limit, max_chars, namespace).
+
+        Returns:
+            BaseMemory: A constructed memory adapter instance.
+
+        Raises:
+            ValueError: If backend is unknown or adapter can\'t be constructed.
+        '''

aip_agents/memory/guidance.pyi

@@ -0,0 +1,3 @@
+from _typeshed import Incomplete
+
+MEM0_MEMORY_RECALL_GUIDANCE: Incomplete

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__.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']