aip-agents-binary 0.5.14__py3-none-any.whl → 0.5.16__py3-none-any.whl

aip_agents/tools/constants.pyi

@@ -0,0 +1,135 @@
+from _typeshed import Incomplete
+from enum import Enum, StrEnum
+
+BOSA_API_BASE_URL: Incomplete
+BOSA_API_KEY: Incomplete
+BOSA_FETCH_MAX_RETRIES: int
+
+class ToolType(StrEnum):
+    """Tool types for BOSA Connector."""
+    GLLM: str
+    LANGCHAIN: str
+
+class Action(Enum):
+    """Actions for BOSA Connector."""
+    GITHUB: str
+    GOOGLE: str
+    GOOGLE_DRIVE: str
+    GOOGLE_MAIL: str
+    TWITTER: str
+
+class GitHubEndpoint(Enum):
+    """GitHub endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    CREATE_ISSUE_HANDLER: str
+    GET_ISSUE_HANDLER: str
+    LIST_ISSUES: str
+    LIST_ISSUES_COMMENTS: str
+    SEARCH_ALL_ISSUES: str
+    GET_COMMITS: str
+    SEARCH_COMMITS: str
+    GET_COLLABORATORS: str
+    GET_RELEASES: str
+    GET_CONTRIBUTORS: str
+    GET_LANGUAGES: str
+    SEARCH_CONTRIBUTIONS: str
+    LIST_PULL_REQUESTS: str
+    SEARCH_PR: str
+    GET_PULL: str
+    GET_ALL_CONTRIBUTOR_COMMIT_ACTIVITIES: str
+    GET_COMMIT_ACTIVITY: str
+    GET_WEEKLY_COMMIT_COUNTS: str
+    GET_USER_CONTRIBUTION_STATISTICS: str
+    LIST_PROJECT_ITEMS: str
+    LIST_PROJECTS: str
+
+class GoogleDriveEndpoint(Enum):
+    """Google Drive endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    SEARCH_FILES: str
+    GET_FILE: str
+    CREATE_FILE: str
+    CREATE_FOLDER: str
+    UPDATE_FILE: str
+    UPDATE_FOLDER: str
+    COPY_FILE: str
+    DELETE_FILE: str
+    SUMMARIZE_FOLDER_FILES_BY_TYPE: str
+    SUMMARIZE_TOTAL_FILES_BY_TYPE: str
+    RECENT_FILES: str
+    CREATE_PERMISSION: str
+    LIST_PERMISSIONS: str
+    GET_PERMISSION: str
+    UPDATE_PERMISSION: str
+    DELETE_PERMISSION: str
+    DOWNLOAD_FILE: str
+
+class GoogleDocsEndpoint(Enum):
+    """Google Docs endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    GET_DOCUMENT: str
+    LIST_DOCUMENTS: str
+    CREATE_DOCUMENT: str
+    UPDATE_DOCUMENT: str
+    COPY_CONTENT: str
+    UPDATE_DOCUMENT_MARKDOWN: str
+    LIST_COMMENTS: str
+    SUMMARIZE_COMMENTS: str
+
+class GoogleEndpoint(Enum):
+    """Google endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    USERINFO: str
+
+class TwitterEndpoint(Enum):
+    """Twitter endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    SEARCH: str
+    GET_TWEETS: str
+    GET_THREAD: str
+    GET_USERS: str
+
+class GoogleMailEndpoint(Enum):
+    """Google Mail endpoints for BOSA Connector."""
+    INTEGRATIONS: str
+    USER_HAS_INTEGRATION: str
+    SUCCESS_AUTHORIZE_CALLBACK: str
+    CREATE_DRAFT: str
+    LIST_DRAFTS: str
+    SEND_DRAFT: str
+    GET_DRAFT: str
+    MODIFY_DRAFT: str
+    LIST_LABELS: str
+    LABEL_STATS: str
+    GET_LABEL_DETAILS: str
+    CREATE_LABELS: str
+    MODIFY_LABELS: str
+    DELETE_LABELS: str
+    SEND_EMAIL: str
+    LIST_EMAILS: str
+    GET_EMAIL_DETAILS: str
+    MODIFY_EMAIL: str
+    DELETE_EMAIL: str
+    TRASH_EMAIL: str
+    UNTRASH_EMAIL: str
+    LIST_THREADS: str
+    THREAD_DETAILS: str
+    MODIFY_THREAD: str
+    GET_AUTO_REPLY: str
+    SET_AUTO_REPLY: str
+    GET_ATTACHMENT: str
+    USERINFO: str
+
+class ActionEndpointMap:
+    """Maps Action enums to their corresponding Endpoint enums."""
+    MAP: dict[Action, type[Enum]]

aip_agents/tools/document_loader/__init__.py

@@ -12,6 +12,8 @@ References:
 
 try:
     from aip_agents.tools.document_loader.base_reader import (  # noqa: F401
+        DOCPROC_AVAILABLE,
+        DOCPROC_MISSING_MESSAGE,
         BaseDocumentReaderTool,
         DocumentReaderInput,
     )
@@ -28,10 +30,15 @@ try:
         "ExcelReaderTool",
         "PDFSplitter",
     ]
+    if not DOCPROC_AVAILABLE:
+        import warnings
+
+        warnings.warn(DOCPROC_MISSING_MESSAGE, ImportWarning)
 except ImportError:
     import warnings
 
     warnings.warn(
-        "Document loader tools not available. Install with: pip install aip-agents[document-loader]", ImportWarning
+        "Document loader tools not available. Install gllm-docproc to enable them.",
+        ImportWarning,
     )
     __all__ = []

aip_agents/tools/document_loader/__init__.pyi

@@ -0,0 +1,7 @@
+from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool as BaseDocumentReaderTool, DocumentReaderInput as DocumentReaderInput
+from aip_agents.tools.document_loader.docx_reader_tool import DocxReaderTool as DocxReaderTool
+from aip_agents.tools.document_loader.excel_reader_tool import ExcelReaderTool as ExcelReaderTool
+from aip_agents.tools.document_loader.pdf_reader_tool import PDFReaderTool as PDFReaderTool
+from aip_agents.tools.document_loader.pdf_splitter import PDFSplitter as PDFSplitter
+
+__all__ = ['BaseDocumentReaderTool', 'DocumentReaderInput', 'PDFReaderTool', 'DocxReaderTool', 'ExcelReaderTool', 'PDFSplitter']

aip_agents/tools/document_loader/base_reader.py

@@ -10,8 +10,8 @@ References:
 import gc
 import os
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
 
-from gllm_docproc.loader.pipeline_loader import PipelineLoader
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, Field
@@ -21,6 +21,46 @@ from aip_agents.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
+DOCPROC_MISSING_MESSAGE = (
+    "gllm-docproc is required for document loader tools but is not installed. "
+    "Install it from your internal registry to enable document processing."
+)
+
+try:
+    from gllm_docproc.loader.pipeline_loader import PipelineLoader
+except ImportError:
+    PipelineLoader = None  # type: ignore[assignment]
+
+if TYPE_CHECKING:
+    from gllm_docproc.loader.pipeline_loader import PipelineLoader as PipelineLoaderType
+else:
+    PipelineLoaderType = Any
+
+
+class _MissingDocprocLoader:
+    """Fallback loader that errors when document processing is attempted."""
+
+    def __init__(self) -> None:
+        self.loaders: list[object] = []
+
+    def add_loader(self, loader: object) -> None:
+        self.loaders.append(loader)
+
+    def load(self, *_args: object, **_kwargs: object) -> list[dict[str, str]]:
+        raise ImportError(DOCPROC_MISSING_MESSAGE)
+
+    def clear_cache(self) -> None:
+        self.loaders.clear()
+
+
+def _build_pipeline_loader() -> "PipelineLoader":
+    if PipelineLoader is None:
+        return _MissingDocprocLoader()
+    return PipelineLoader()  # type: ignore[misc]
+
+
+DOCPROC_AVAILABLE = PipelineLoader is not None
+
 
 class BaseDocumentConfig(BaseModel):
     """Base tool configuration schema for document processing with batching functionality.
@@ -78,7 +118,7 @@ class BaseDocumentReaderTool(BaseTool, ABC):
     description: str = "Read a document file and extract its text content."
     args_schema: type[BaseModel] = DocumentReaderInput
     tool_config_schema: type[BaseModel] = BaseDocumentConfig
-    loader: PipelineLoader = Field(default_factory=PipelineLoader)
+    loader: PipelineLoaderType = Field(default_factory=_build_pipeline_loader)
 
     def __init__(self):
         """Initialize the base document reader tool."""

aip_agents/tools/document_loader/base_reader.pyi

@@ -0,0 +1,75 @@
+from _typeshed import Incomplete
+from abc import ABC
+from aip_agents.utils.logger import get_logger as get_logger
+from gllm_docproc.loader.pipeline_loader import PipelineLoader as PipelineLoaderType
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel
+from typing import Any
+
+logger: Incomplete
+DOCPROC_MISSING_MESSAGE: str
+PipelineLoaderType = Any
+
+class _MissingDocprocLoader:
+    """Fallback loader that errors when document processing is attempted."""
+    loaders: list[object]
+    def __init__(self) -> None: ...
+    def add_loader(self, loader: object) -> None: ...
+    def load(self, *_args: object, **_kwargs: object) -> list[dict[str, str]]: ...
+    def clear_cache(self) -> None: ...
+
+DOCPROC_AVAILABLE: Incomplete
+
+class BaseDocumentConfig(BaseModel):
+    """Base tool configuration schema for document processing with batching functionality.
+
+    This configuration enables page-by-page batching to optimize memory usage when
+    processing large document files. When batching is enabled, documents are processed
+    sequentially by pages rather than loading the entire document into memory at once.
+
+    Attributes:
+        batching (bool): Enable page-by-page batching to reduce memory usage.
+            When True, documents are processed page by page sequentially.
+            When False, maintains current behavior of loading entire document.
+            Defaults to False for backward compatibility.
+        batch_size (int): Number of pages to process in each batch.
+            Must be between 1 and 100 pages inclusive.
+            Larger batch sizes may use more memory but could be more efficient.
+            Smaller batch sizes use less memory but may have more overhead.
+            Defaults to 10 for balanced memory usage and efficiency.
+
+    Examples:
+        >>> # Default configuration (no batching)
+        >>> config = BaseDocumentConfig()
+        >>> print(config.batching)  # False
+        >>> print(config.batch_size)  # 10
+
+        >>> # Enable batching with single page processing
+        >>> config = BaseDocumentConfig(batching=True, batch_size=1)
+
+        >>> # Enable batching with multi-page batches
+        >>> config = BaseDocumentConfig(batching=True, batch_size=3)
+    """
+    batching: bool
+    batch_size: int
+
+class DocumentReaderInput(BaseModel):
+    """Input schema for the DocumentReader tool."""
+    file_path: str
+
+class BaseDocumentReaderTool(BaseTool, ABC):
+    """Base tool to read and extract text from document files."""
+    name: str
+    description: str
+    args_schema: type[BaseModel]
+    tool_config_schema: type[BaseModel]
+    loader: PipelineLoaderType
+    def __init__(self) -> None:
+        """Initialize the base document reader tool."""
+    def cleanup_memory(self) -> None:
+        """Explicitly clean up memory and force garbage collection.
+
+        This method can be called after processing to minimize memory usage.
+        While it won't reset memory to exactly 0, it will free up as much
+        memory as possible by clearing internal caches and forcing garbage collection.
+        """

aip_agents/tools/document_loader/docx_reader_tool.py

@@ -8,9 +8,19 @@ References:
       reader/docx_reader_tool.py
 """
 
-from gllm_docproc.loader.docx import DOCX2PythonLoader, PythonDOCXTableLoader
+try:
+    from gllm_docproc.loader.docx import DOCX2PythonLoader, PythonDOCXTableLoader
+except ImportError as exc:
+    DOCX2PythonLoader = None  # type: ignore[assignment]
+    PythonDOCXTableLoader = None  # type: ignore[assignment]
+    _DOCPROC_IMPORT_ERROR: Exception | None = exc
+else:
+    _DOCPROC_IMPORT_ERROR = None
 
-from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool
+from aip_agents.tools.document_loader.base_reader import DOCPROC_MISSING_MESSAGE, BaseDocumentReaderTool
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class DocxReaderTool(BaseDocumentReaderTool):
@@ -20,6 +30,11 @@ class DocxReaderTool(BaseDocumentReaderTool):
     description: str = "Read a Word document and extract its text content. Input should be the path to the Word file."
 
     def _setup_loader(self):
+        if DOCX2PythonLoader is None or PythonDOCXTableLoader is None:
+            logger.warning(DOCPROC_MISSING_MESSAGE)
+            if _DOCPROC_IMPORT_ERROR is not None:
+                logger.debug("gllm_docproc import failed: %s", _DOCPROC_IMPORT_ERROR)
+            return
         self.loader.add_loader(DOCX2PythonLoader())
         self.loader.add_loader(PythonDOCXTableLoader())
 

aip_agents/tools/document_loader/docx_reader_tool.pyi

@@ -0,0 +1,10 @@
+from _typeshed import Incomplete
+from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool as BaseDocumentReaderTool, DOCPROC_MISSING_MESSAGE as DOCPROC_MISSING_MESSAGE
+from aip_agents.utils.logger import get_logger as get_logger
+
+logger: Incomplete
+
+class DocxReaderTool(BaseDocumentReaderTool):
+    """Tool to read and extract text from Word documents."""
+    name: str
+    description: str

aip_agents/tools/document_loader/excel_reader_tool.py

@@ -10,9 +10,15 @@ Authors:
 import zipfile
 from pathlib import Path
 
-from gllm_docproc.loader.xlsx import OpenpyxlLoader
-
-from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool
+try:
+    from gllm_docproc.loader.xlsx import OpenpyxlLoader
+except ImportError as exc:
+    OpenpyxlLoader = None  # type: ignore[assignment]
+    _DOCPROC_IMPORT_ERROR: Exception | None = exc
+else:
+    _DOCPROC_IMPORT_ERROR = None
+
+from aip_agents.tools.document_loader.base_reader import DOCPROC_MISSING_MESSAGE, BaseDocumentReaderTool
 from aip_agents.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -50,6 +56,11 @@ class ExcelReaderTool(BaseDocumentReaderTool):
         This method initializes the OpenpyxlLoader which handles extraction
         of content from Excel files and formatting as Markdown tables.
         """
+        if OpenpyxlLoader is None:
+            logger.warning(DOCPROC_MISSING_MESSAGE)
+            if _DOCPROC_IMPORT_ERROR is not None:
+                logger.debug("gllm_docproc import failed: %s", _DOCPROC_IMPORT_ERROR)
+            return
         self.loader.add_loader(OpenpyxlLoader())
 
     def _run_with_batching(self, file_path: str, batch_size: int) -> str:

aip_agents/tools/document_loader/excel_reader_tool.pyi

@@ -0,0 +1,26 @@
+from _typeshed import Incomplete
+from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool as BaseDocumentReaderTool, DOCPROC_MISSING_MESSAGE as DOCPROC_MISSING_MESSAGE
+from aip_agents.utils.logger import get_logger as get_logger
+
+logger: Incomplete
+
+class ExcelReaderTool(BaseDocumentReaderTool):
+    '''Tool to read and extract content from Excel files.
+
+    This tool reads Excel files (.xlsx, .xlsm) and extracts their content using
+    the gllm_docproc loader pipeline. The content is formatted as Markdown tables
+    for easy readability.
+
+    Features:
+    - Supports .xlsx and .xlsm formats
+    - Extracts all sheets or specific sheets
+    - Formats output as Markdown tables
+    - Configurable row limits and file size limits
+
+    Examples:
+        >>> tool = ExcelReaderTool()
+        >>> result = tool._run("/tmp/data.xlsx")
+        >>> print(result)
+    '''
+    name: str
+    description: str

aip_agents/tools/document_loader/pdf_reader_tool.py

@@ -8,9 +8,16 @@ References:
       reader/pdf_reader_tool.py
 """
 
-from gllm_docproc.loader.pdf import PDFPlumberLoader, PyMuPDFLoader
-
-from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool
+try:
+    from gllm_docproc.loader.pdf import PDFPlumberLoader, PyMuPDFLoader
+except ImportError as exc:
+    PDFPlumberLoader = None  # type: ignore[assignment]
+    PyMuPDFLoader = None  # type: ignore[assignment]
+    _DOCPROC_IMPORT_ERROR: Exception | None = exc
+else:
+    _DOCPROC_IMPORT_ERROR = None
+
+from aip_agents.tools.document_loader.base_reader import DOCPROC_MISSING_MESSAGE, BaseDocumentReaderTool
 from aip_agents.tools.document_loader.pdf_splitter import PDFSplitter
 from aip_agents.utils.logger import get_logger
 
@@ -24,6 +31,11 @@ class PDFReaderTool(BaseDocumentReaderTool):
     description: str = "Read a PDF file and extract its text content. Input should be the path to the PDF file."
 
     def _setup_loader(self):
+        if PDFPlumberLoader is None or PyMuPDFLoader is None:
+            logger.warning(DOCPROC_MISSING_MESSAGE)
+            if _DOCPROC_IMPORT_ERROR is not None:
+                logger.debug("gllm_docproc import failed: %s", _DOCPROC_IMPORT_ERROR)
+            return
         self.loader.add_loader(PyMuPDFLoader())
         self.loader.add_loader(PDFPlumberLoader())
 
@@ -43,7 +55,7 @@ class PDFReaderTool(BaseDocumentReaderTool):
         Raises:
             FileNotFoundError: If the input PDF file doesn't exist
             ValueError: If batch_size is invalid or PDF processing fails
-            Exception: For other unexpected errors during PDF splitting
+            RuntimeError: For other unexpected errors during PDF splitting
         """
         logger.info(f"Splitting PDF file '{file_path}' with batch_size={batch_size}")
 
@@ -64,4 +76,4 @@ class PDFReaderTool(BaseDocumentReaderTool):
             raise
         except Exception as e:
             logger.error(f"Unexpected error splitting PDF '{file_path}': {str(e)}")
-            raise Exception(f"Failed to split PDF file: {str(e)}") from e
+            raise RuntimeError(f"Failed to split PDF file: {str(e)}") from e

aip_agents/tools/document_loader/pdf_reader_tool.pyi

@@ -0,0 +1,11 @@
+from _typeshed import Incomplete
+from aip_agents.tools.document_loader.base_reader import BaseDocumentReaderTool as BaseDocumentReaderTool, DOCPROC_MISSING_MESSAGE as DOCPROC_MISSING_MESSAGE
+from aip_agents.tools.document_loader.pdf_splitter import PDFSplitter as PDFSplitter
+from aip_agents.utils.logger import get_logger as get_logger
+
+logger: Incomplete
+
+class PDFReaderTool(BaseDocumentReaderTool):
+    """Tool to read and extract text from PDF files."""
+    name: str
+    description: str

aip_agents/tools/document_loader/pdf_splitter.pyi

@@ -0,0 +1,18 @@
+class PDFSplitter:
+    """Utility class for splitting PDF files into page-based temporary files."""
+    @staticmethod
+    def split_by_pages(file_path: str, batch_size: int) -> list[str]:
+        """Split PDF into temporary files containing batch_size pages each.
+
+        Args:
+            file_path: Path to the PDF file to split
+            batch_size: Number of pages to include in each batch (1-10)
+
+        Returns:
+            List of temporary file paths containing the split PDF batches
+
+        Raises:
+            FileNotFoundError: If the input PDF file doesn't exist
+            ValueError: If batch_size is invalid or PDF is corrupted
+            Exception: For other PDF processing errors
+        """

aip_agents/tools/gl_connector/__init__.pyi

@@ -0,0 +1,3 @@
+from aip_agents.tools.gl_connector.tool import GLConnectorTool as GLConnectorTool
+
+__all__ = ['GLConnectorTool']

aip_agents/tools/gl_connector/tool.pyi

@@ -0,0 +1,74 @@
+from _typeshed import Incomplete
+from aip_agents.tools.constants import ToolType as ToolType
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool
+from typing import Any
+
+class _InjectedTool(BaseTool):
+    """Wrap a BaseTool to inject token and optional identifier into inputs."""
+    model_config: Incomplete
+    def __init__(self, base_tool: BaseTool, token: str, identifier: str | None) -> None:
+        """Initialize the injected tool wrapper.
+
+        Args:
+            base_tool: The base tool to wrap.
+            token: Authentication token to inject into tool inputs.
+            identifier: Optional identifier to inject into tool inputs.
+
+        Returns:
+            None
+        """
+    def invoke(self, input: Any, config: RunnableConfig | None = None, **kwargs: Any) -> Any:
+        """Invoke the tool with token and optional identifier injected.
+
+        Args:
+            input: Tool input to process.
+            config: Optional runnable configuration.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The result of invoking the tool with injected parameters.
+        """
+    async def ainvoke(self, input: Any, config: RunnableConfig | None = None, **kwargs: Any) -> Any:
+        """Invoke the tool asynchronously with token and optional identifier injected.
+
+        Args:
+            input: Tool input to process.
+            config: Optional runnable configuration.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The result of invoking the tool with injected parameters.
+        """
+    def run(self, tool_input: Any, **kwargs: Any) -> Any:
+        """Run the tool with token and optional identifier injected.
+
+        Args:
+            tool_input: Tool input to process.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The result of running the tool with injected parameters.
+        """
+    async def arun(self, tool_input: Any, **kwargs: Any) -> Any:
+        """Run the tool asynchronously with token and optional identifier injected.
+
+        Args:
+            tool_input: Tool input to process.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The result of running the tool with injected parameters.
+        """
+
+def GLConnectorTool(tool_name: str, *, api_key: str | None = None, identifier: str | None = None) -> BaseTool:
+    """Create a single GL Connector tool by exact tool name.
+
+    Args:
+        tool_name: Exact tool name (not module name).
+        api_key: Optional override for BOSA API key.
+        identifier: Optional override for BOSA identifier.
+
+    Returns:
+        A single LangChain BaseTool with token injection.
+    """

aip_agents/tools/memory_search/__init__.pyi

@@ -0,0 +1,5 @@
+from aip_agents.tools.memory_search.base import LongTermMemorySearchTool as LongTermMemorySearchTool
+from aip_agents.tools.memory_search.mem0 import MEMORY_SEARCH_TOOL_NAME as MEMORY_SEARCH_TOOL_NAME, Mem0SearchInput as Mem0SearchInput, Mem0SearchTool as Mem0SearchTool
+from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput as LongTermMemorySearchInput, MemoryConfig as MemoryConfig
+
+__all__ = ['MemoryConfig', 'LongTermMemorySearchInput', 'LongTermMemorySearchTool', 'Mem0SearchInput', 'Mem0SearchTool', 'MEMORY_SEARCH_TOOL_NAME']

aip_agents/tools/memory_search/base.pyi

@@ -0,0 +1,69 @@
+from _typeshed import Incomplete
+from abc import ABC
+from aip_agents.memory.constants import MemoryDefaults as MemoryDefaults
+from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput as LongTermMemorySearchInput, MemoryConfig as MemoryConfig
+from aip_agents.utils.logger import get_logger as get_logger
+from collections.abc import Callable
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel as BaseModel
+from typing import Any, ClassVar, Protocol
+
+logger: Incomplete
+
+class LongTermMemoryBackend(Protocol):
+    """Protocol for memory adapters that support retrieval and formatting."""
+    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 matching the given query and filters.
+
+        Args:
+            query: Optional search query string.
+            user_id: User identifier for scoped retrieval.
+            limit: Maximum number of results to return.
+            filters: Optional filters to apply to the retrieval.
+            page: Page number for pagination.
+
+        Returns:
+            List of memory hit dictionaries.
+        """
+    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 output.
+            with_tag: Whether to wrap output with memory tags.
+
+        Returns:
+            Formatted string representation of memory hits.
+        """
+
+class LongTermMemorySearchTool(BaseTool, ABC):
+    """Abstract base class for provider-specific long-term memory search tools."""
+    name: str
+    description: str
+    args_schema: type[LongTermMemorySearchInput]
+    tool_config_schema: type[BaseModel]
+    memory: LongTermMemoryBackend
+    default_user_id: str | None
+    user_id_provider: Callable[[], str | None] | None
+    MINIMUM_MEMORY_RETRIEVAL: ClassVar[int]
+    LOG_PREFIX: ClassVar[str]
+    def __init__(self, memory: LongTermMemoryBackend, *, default_user_id: str | None = None, user_id_provider: Callable[[], str | None] | None = None, **kwargs: Any) -> None:
+        """Initialize the long-term memory search tool.
+
+        Args:
+            memory: Memory adapter instance with retrieve() and format_hits() methods.
+            default_user_id: Default user ID to use if not provided in metadata.
+            user_id_provider: Callable that returns a user ID.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+    def format_hits(self, hits: list[dict[str, Any]], with_tag: bool = False) -> str:
+        """Format hits into a string with optional tags.
+
+        Args:
+            hits: List of memory hit dictionaries to format.
+            with_tag: Whether to wrap the output with memory tags.
+
+        Returns:
+            Formatted string representation of the memory hits.
+        """

aip_agents/tools/memory_search/mem0.pyi

@@ -0,0 +1,19 @@
+from _typeshed import Incomplete
+from aip_agents.memory.constants import MemoryDefaults as MemoryDefaults
+from aip_agents.tools.memory_search.base import LongTermMemorySearchTool as LongTermMemorySearchTool
+from aip_agents.tools.memory_search.schema import LongTermMemorySearchInput as LongTermMemorySearchInput
+from aip_agents.utils.datetime import is_valid_date_string as is_valid_date_string, next_day_iso as next_day_iso
+from aip_agents.utils.logger import get_logger as get_logger
+from typing import ClassVar
+
+logger: Incomplete
+MEMORY_SEARCH_TOOL_NAME: str
+
+class Mem0SearchTool(LongTermMemorySearchTool):
+    """Mem0-specific implementation of the long-term memory search tool."""
+    name: str
+    description: str
+    args_schema: type[LongTermMemorySearchInput]
+    LOG_PREFIX: ClassVar[str]
+    METADATA_FILTER_BLOCKLIST: ClassVar[set[str]]
+Mem0SearchInput = LongTermMemorySearchInput