ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/tools/o365/base.py

@@ -0,0 +1,295 @@
+"""
+Office365 Tools for AI-Parrot.
+
+Base classes and utilities for Office365 integration tools.
+Supports multiple authentication modes:
+- Direct (admin/client credentials)
+- On-Behalf-Of (OBO)
+- Delegated (interactive user login)
+"""
+from typing import Dict, Any, Optional, Type, List
+from abc import abstractmethod
+import contextlib
+from datetime import datetime
+import asyncio
+from pydantic import BaseModel, Field
+from navconfig.logging import logging
+from ..abstract import AbstractTool, AbstractToolArgsSchema, ToolResult
+# Import your O365Client
+from ...interfaces.o365 import O365Client
+
+
+class O365AuthMode:
+    """Authentication modes for Office365 tools."""
+    DIRECT = "direct"  # Admin/App credentials
+    OBO = "on_behalf_of"  # On-Behalf-Of flow
+    DELEGATED = "delegated"  # Interactive user login
+    CACHED = "cached"  # Use cached interactive session
+
+
+class O365ToolArgsSchema(AbstractToolArgsSchema):
+    """Base schema for Office365 tool arguments."""
+    auth_mode: Optional[str] = Field(
+        default=None,
+        description="Authentication mode: 'direct', 'on_behalf_of', 'delegated', or 'cached'. If not provided, uses the tool's default_auth_mode."
+    )
+    user_assertion: Optional[str] = Field(
+        default=None,
+        description="User assertion token for OBO flow"
+    )
+    user_id: Optional[str] = Field(
+        default=None,
+        description=(
+            "Target mailbox user principal name or ID. Required when using app-only "
+            "(client credentials) authentication to access user-specific resources."
+        )
+    )
+
+
+class O365Tool(AbstractTool):
+    """
+    Base class for Office365 tools that interact with Microsoft Graph API.
+
+    This class provides:
+    - Integration with O365Client
+    - Multiple authentication modes
+    - Error handling and logging
+    - Async execution support
+
+    Subclasses should implement:
+    - _execute_graph_operation(): Perform the actual Graph API operation
+
+    Authentication Modes:
+    1. DIRECT: Uses client credentials (app-only access)
+       - Best for: Admin operations, bulk operations
+       - Requires: client_id, client_secret, tenant_id
+
+    2. ON_BEHALF_OF: Uses OBO flow with user assertion
+       - Best for: Acting on behalf of authenticated user
+       - Requires: client_id, client_secret, tenant_id, user_assertion
+
+    3. DELEGATED: Interactive user login
+       - Best for: User-specific operations with full permissions
+       - Requires: Interactive browser login
+
+    4. CACHED: Reuse cached interactive session
+       - Best for: Subsequent operations after interactive login
+       - Requires: Previous interactive_login() call
+    """
+
+    name: str = "o365_base"
+    description: str = "Base Office365 tool"
+    args_schema: Type[BaseModel] = O365ToolArgsSchema
+    return_direct: bool = False
+
+    def __init__(
+        self,
+        credentials: Optional[Dict[str, Any]] = None,
+        default_auth_mode: str = O365AuthMode.DIRECT,
+        scopes: Optional[List[str]] = None,
+        **kwargs
+    ):
+        """
+        Initialize Office365 tool.
+
+        Args:
+            credentials: Dict with O365 credentials (client_id, client_secret, tenant_id, etc.)
+            default_auth_mode: Default authentication mode to use
+            scopes: List of Graph API scopes
+            **kwargs: Additional arguments for AbstractTool
+        """
+        super().__init__(**kwargs)
+
+        self.credentials = credentials or {}
+        self.default_auth_mode = default_auth_mode
+        self.scopes = scopes
+        self._client: Optional[O365Client] = None
+        self._client_cache: Dict[str, O365Client] = {}  # Cache clients by auth mode
+        self._authenticated = False  # Track auth state
+        self.logger = logging.getLogger(f'Parrot.Tools.{self.__class__.__name__}')
+
+    async def _get_client(
+        self,
+        auth_mode: str = None,
+        user_assertion: str = None,
+        scopes: List[str] = None
+    ) -> O365Client:
+        """
+        Get or create O365Client based on authentication mode.
+
+        Args:
+            auth_mode: Authentication mode to use
+            user_assertion: User assertion for OBO flow
+            scopes: Graph API scopes
+
+        Returns:
+            Configured O365Client instance
+        """
+        auth_mode = auth_mode or self.default_auth_mode
+        scopes = scopes or self.scopes
+        cache_key = f"{auth_mode}_{user_assertion or 'none'}"
+
+        # Check cache
+        if cache_key in self._client_cache:
+            cached_client = self._client_cache[cache_key]
+            cached_client.set_auth_mode(auth_mode)
+            return cached_client
+
+        # Create new client
+        client_credentials = self.credentials.copy()
+
+        # Handle OBO flow
+        if auth_mode == O365AuthMode.OBO:
+            if not user_assertion:
+                raise ValueError("user_assertion is required for OBO authentication")
+            client_credentials['assertion'] = user_assertion
+
+        # Create and configure client
+        client = O365Client(
+            credentials=client_credentials
+        )
+        client.processing_credentials()
+        client.set_auth_mode(auth_mode)
+
+        try:
+            if auth_mode == O365AuthMode.DIRECT:
+                # Client credentials flow
+                self.logger.info("Using direct (client credentials) authentication")
+                await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    client.acquire_token,
+                    scopes or self.scopes
+                )
+
+            elif auth_mode == O365AuthMode.OBO:
+                # On-Behalf-Of flow
+                self.logger.info("Using On-Behalf-Of authentication")
+                await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    client.acquire_token_on_behalf_of,
+                    user_assertion,
+                    scopes or self.scopes
+                )
+
+            elif auth_mode == O365AuthMode.DELEGATED:
+                if not self._authenticated:
+                    await client.interactive_login(scopes=scopes)
+                    self._authenticated = True
+                else:
+                    await client.ensure_interactive_session(scopes=scopes)
+            elif auth_mode == O365AuthMode.CACHED:
+                # Use cached session
+                self.logger.info("Using cached interactive session")
+                await client.ensure_interactive_session(scopes=scopes or self.scopes)
+
+            else:
+                raise ValueError(
+                    f"Unknown authentication mode: {auth_mode}"
+                )
+
+            # Cache the client
+            self._client_cache[cache_key] = client
+
+            return client
+
+        except Exception as e:
+            self.logger.error(f"Failed to authenticate O365Client: {e}")
+            raise
+
+    @abstractmethod
+    async def _execute_graph_operation(
+        self,
+        client: O365Client,
+        **kwargs
+    ) -> Any:
+        """
+        Execute the actual Graph API operation.
+
+        Subclasses must implement this method to perform specific operations.
+
+        Args:
+            client: Authenticated O365Client instance
+            **kwargs: Operation-specific parameters
+
+        Returns:
+            Operation result
+        """
+        pass
+
+    async def _execute(self, **kwargs) -> ToolResult:
+        """
+        Execute the Office365 tool.
+
+        This method handles:
+        1. Authentication setup
+        2. Calling the specific Graph operation
+        3. Error handling
+        4. Result formatting
+
+        Args:
+            **kwargs: Tool arguments including auth parameters
+
+        Returns:
+            ToolResult with operation outcome
+        """
+        start_time = datetime.now()
+
+        try:
+            # Extract auth parameters
+            auth_mode = kwargs.pop('auth_mode', None) or self.default_auth_mode
+            user_assertion = kwargs.pop('user_assertion', None)
+            scopes = kwargs.pop('scopes', None) or self.scopes
+
+            self.logger.debug(
+                f"Executing with auth_mode={auth_mode}, "
+                f"default_auth_mode={self.default_auth_mode}"
+            )
+
+            # Get authenticated client
+            client = await self._get_client(
+                auth_mode=auth_mode,
+                user_assertion=user_assertion,
+                scopes=scopes
+            )
+
+            # Execute the operation
+            self.logger.info(f"Executing {self.name} operation")
+            result = await self._execute_graph_operation(client, **kwargs)
+
+            # Calculate execution time
+            execution_time = (datetime.now() - start_time).total_seconds()
+
+            return ToolResult(
+                status="success",
+                result=result,
+                metadata={
+                    "tool": self.name,
+                    "auth_mode": auth_mode,
+                    "execution_time": execution_time,
+                    "timestamp": datetime.now().isoformat()
+                }
+            )
+
+        except Exception as e:
+            self.logger.error(f"Error executing {self.name}: {e}", exc_info=True)
+            execution_time = (datetime.now() - start_time).total_seconds()
+
+            return ToolResult(
+                status="error",
+                result=None,
+                error=str(e),
+                metadata={
+                    "tool": self.name,
+                    "execution_time": execution_time,
+                    "timestamp": datetime.now().isoformat()
+                }
+            )
+
+    async def cleanup(self):
+        """Clean up resources."""
+        for client in self._client_cache.values():
+            try:
+                await client.close()
+            except Exception as e:
+                self.logger.warning(f"Error closing client: {e}")
+        self._client_cache.clear()