pltr-cli 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

pltr/services/admin.py

@@ -3,10 +3,12 @@ Admin service wrapper for Foundry SDK admin operations.
 Provides a high-level interface for user, group, role, and organization management.
 """
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Callable
 import json
 
 from .base import BaseService
+from ..utils.pagination import PaginationConfig, PaginationResult
+from ..config.settings import Settings
 
 
 class AdminService(BaseService):
@@ -23,6 +25,8 @@ class AdminService(BaseService):
         """
         List all users in the organization.
 
+        DEPRECATED: Use list_users_paginated() instead for better pagination support.
+
         Args:
             page_size: Maximum number of users to return per page
             page_token: Token for pagination (from previous response)
@@ -38,6 +42,47 @@ class AdminService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to list users: {str(e)}")
 
+    def list_users_paginated(
+        self,
+        config: PaginationConfig,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> PaginationResult:
+        """
+        List users with full pagination control.
+
+        Args:
+            config: Pagination configuration (page_size, max_pages, etc.)
+            progress_callback: Optional callback(page_num, items_count)
+
+        Returns:
+            PaginationResult with users and metadata
+
+        Example:
+            >>> config = PaginationConfig(page_size=50, max_pages=2)
+            >>> result = service.list_users_paginated(config)
+            >>> print(f"Fetched {result.metadata.items_fetched} users")
+        """
+        try:
+            settings = Settings()
+
+            def fetch_page(page_token: Optional[str]) -> Dict[str, Any]:
+                """Fetch a single page of users."""
+                iterator = self.service.User.list(
+                    page_size=config.page_size or settings.get("page_size", 20),
+                    page_token=page_token,
+                )
+                # ResourceIterator has .data and .next_page_token attributes
+                # Extract them properly for the pagination handler
+                return {
+                    "data": [self._serialize_response(user) for user in iterator.data],
+                    "next_page_token": iterator.next_page_token,
+                }
+
+            # Use response pagination handler
+            return self._paginate_response(fetch_page, config, progress_callback)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list users: {str(e)}")
+
     def get_user(self, user_id: str) -> Dict[str, Any]:
         """
         Get a specific user by ID.
@@ -127,6 +172,43 @@ class AdminService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to revoke tokens for user {user_id}: {str(e)}")
 
+    def delete_user(self, user_id: str) -> Dict[str, Any]:
+        """
+        Delete a specific user.
+
+        Args:
+            user_id: The user ID or RID
+
+        Returns:
+            Dictionary containing operation result
+        """
+        try:
+            self.service.User.delete(user_id)
+            return {
+                "success": True,
+                "message": f"User {user_id} deleted successfully",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to delete user {user_id}: {str(e)}")
+
+    def get_batch_users(self, user_ids: List[str]) -> Dict[str, Any]:
+        """
+        Batch retrieve multiple users (up to 500).
+
+        Args:
+            user_ids: List of user IDs or RIDs
+
+        Returns:
+            Dictionary containing user information
+        """
+        if len(user_ids) > 500:
+            raise ValueError("Maximum batch size is 500 users")
+        try:
+            response = self.service.User.get_batch(body=user_ids)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get users batch: {str(e)}")
+
     # Group Management Methods
     def list_groups(
         self, page_size: Optional[int] = None, page_token: Optional[str] = None
@@ -142,10 +224,16 @@ class AdminService(BaseService):
             Dictionary containing group list and pagination info
         """
         try:
-            response = self.service.Group.list(
+            iterator = self.service.Group.list(
                 page_size=page_size, page_token=page_token
             )
-            return self._serialize_response(response)
+            # ResourceIterator has .data and .next_page_token attributes
+            # Return structure compatible with formatter.display()
+            groups = [self._serialize_response(group) for group in iterator.data]
+            return {
+                "data": groups,
+                "next_page_token": iterator.next_page_token,
+            }
         except Exception as e:
             raise RuntimeError(f"Failed to list groups: {str(e)}")
 
@@ -239,6 +327,136 @@ class AdminService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to delete group {group_id}: {str(e)}")
 
+    def get_batch_groups(self, group_ids: List[str]) -> Dict[str, Any]:
+        """
+        Batch retrieve multiple groups (up to 500).
+
+        Args:
+            group_ids: List of group IDs or RIDs
+
+        Returns:
+            Dictionary containing group information
+        """
+        if len(group_ids) > 500:
+            raise ValueError("Maximum batch size is 500 groups")
+        try:
+            response = self.service.Group.get_batch(body=group_ids)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get groups batch: {str(e)}")
+
+    # Marking Management Methods
+    def list_markings(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        List all markings.
+
+        Args:
+            page_size: Maximum number of markings to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing marking list and pagination info
+        """
+        try:
+            response = self.service.Marking.list(
+                page_size=page_size, page_token=page_token, preview=True
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list markings: {str(e)}")
+
+    def get_marking(self, marking_id: str) -> Dict[str, Any]:
+        """
+        Get a specific marking by ID.
+
+        Args:
+            marking_id: The marking ID
+
+        Returns:
+            Dictionary containing marking information
+        """
+        try:
+            response = self.service.Marking.get(marking_id, preview=True)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get marking {marking_id}: {str(e)}")
+
+    def get_batch_markings(self, marking_ids: List[str]) -> Dict[str, Any]:
+        """
+        Batch retrieve multiple markings (up to 500).
+
+        Args:
+            marking_ids: List of marking IDs
+
+        Returns:
+            Dictionary containing marking information
+        """
+        if len(marking_ids) > 500:
+            raise ValueError("Maximum batch size is 500 markings")
+        try:
+            response = self.service.Marking.get_batch(body=marking_ids, preview=True)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get markings batch: {str(e)}")
+
+    def create_marking(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        category_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Create a new marking.
+
+        Args:
+            name: The marking name
+            description: Optional marking description
+            category_id: Optional category ID for the marking
+
+        Returns:
+            Dictionary containing created marking information
+        """
+        try:
+            create_params: Dict[str, Any] = {"name": name, "preview": True}
+            if description:
+                create_params["description"] = description
+            if category_id:
+                create_params["category_id"] = category_id
+
+            response = self.service.Marking.create(**create_params)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create marking '{name}': {str(e)}")
+
+    def replace_marking(
+        self,
+        marking_id: str,
+        name: str,
+        description: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Replace/update an existing marking.
+
+        Args:
+            marking_id: The marking ID
+            name: The new marking name
+            description: Optional new marking description
+
+        Returns:
+            Dictionary containing updated marking information
+        """
+        try:
+            replace_params: Dict[str, Any] = {"name": name, "preview": True}
+            if description:
+                replace_params["description"] = description
+
+            response = self.service.Marking.replace(marking_id, **replace_params)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to replace marking {marking_id}: {str(e)}")
+
     # Organization Management Methods
     def get_organization(self, organization_id: str) -> Dict[str, Any]:
         """
@@ -258,6 +476,98 @@ class AdminService(BaseService):
                 f"Failed to get organization {organization_id}: {str(e)}"
             )
 
+    def create_organization(
+        self,
+        name: str,
+        enrollment_rid: str,
+        admin_ids: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Create a new organization.
+
+        Args:
+            name: The organization name
+            enrollment_rid: The enrollment RID
+            admin_ids: Optional list of admin user IDs
+
+        Returns:
+            Dictionary containing created organization information
+        """
+        try:
+            create_params: Dict[str, Any] = {
+                "name": name,
+                "enrollment_rid": enrollment_rid,
+                "preview": True,
+            }
+            if admin_ids:
+                create_params["admin_ids"] = admin_ids
+
+            response = self.service.Organization.create(**create_params)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create organization '{name}': {str(e)}")
+
+    def replace_organization(
+        self,
+        organization_rid: str,
+        name: str,
+        description: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Replace/update an existing organization.
+
+        Args:
+            organization_rid: The organization RID
+            name: The new organization name
+            description: Optional new organization description
+
+        Returns:
+            Dictionary containing updated organization information
+        """
+        try:
+            replace_params: Dict[str, Any] = {"name": name, "preview": True}
+            if description:
+                replace_params["description"] = description
+
+            response = self.service.Organization.replace(
+                organization_rid, **replace_params
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to replace organization {organization_rid}: {str(e)}"
+            )
+
+    def list_available_roles(
+        self,
+        organization_rid: str,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        List available roles for an organization.
+
+        Args:
+            organization_rid: The organization RID
+            page_size: Maximum number of roles to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing role list and pagination info
+        """
+        try:
+            response = self.service.Organization.list_available_roles(
+                organization_rid,
+                page_size=page_size,
+                page_token=page_token,
+                preview=True,
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to list available roles for organization {organization_rid}: {str(e)}"
+            )
+
     # Role Management Methods
     def get_role(self, role_id: str) -> Dict[str, Any]:
         """
@@ -275,6 +585,24 @@ class AdminService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to get role {role_id}: {str(e)}")
 
+    def get_batch_roles(self, role_ids: List[str]) -> Dict[str, Any]:
+        """
+        Batch retrieve multiple roles (up to 500).
+
+        Args:
+            role_ids: List of role IDs or RIDs
+
+        Returns:
+            Dictionary containing role information
+        """
+        if len(role_ids) > 500:
+            raise ValueError("Maximum batch size is 500 roles")
+        try:
+            response = self.service.Role.get_batch(body=role_ids, preview=True)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get roles batch: {str(e)}")
+
     def _serialize_response(self, response: Any) -> Dict[str, Any]:
         """
         Convert response object to serializable dictionary.

pltr/services/aip_agents.py

@@ -0,0 +1,147 @@
+"""
+AIP Agents service wrapper for Foundry SDK.
+Provides access to AIP Agent operations including agent details, sessions, and versions.
+"""
+
+from typing import Any, Dict, Optional
+from .base import BaseService
+from ..utils.pagination import PaginationConfig, PaginationResult
+
+
+class AipAgentsService(BaseService):
+    """Service wrapper for Foundry AIP Agents operations."""
+
+    def _get_service(self) -> Any:
+        """Get the Foundry AIP agents service."""
+        return self.client.aip_agents
+
+    # ===== Agent Operations =====
+
+    def get_agent(
+        self, agent_rid: str, version: Optional[str] = None, preview: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Get details for an AIP Agent.
+
+        Args:
+            agent_rid: Agent Resource Identifier
+                Format: ri.foundry.main.agent.<id>
+            version: Optional agent version string (e.g., "1.0")
+                If not specified, returns latest published version
+            preview: Enable preview mode (default: True)
+
+        Returns:
+            Agent information dictionary containing:
+            - rid: Agent resource identifier
+            - version: Agent version
+            - metadata: Agent metadata (displayName, description, etc.)
+            - parameters: Application variables/parameters
+
+        Raises:
+            RuntimeError: If the operation fails
+
+        Example:
+            >>> service = AipAgentsService()
+            >>> agent = service.get_agent("ri.foundry.main.agent.abc123")
+            >>> print(agent['metadata']['displayName'])
+        """
+        try:
+            agent = self.service.Agent.get(agent_rid, version=version, preview=preview)
+            return self._serialize_response(agent)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get agent {agent_rid}: {e}")
+
+    # ===== Session Operations =====
+
+    def list_sessions(
+        self, agent_rid: str, config: PaginationConfig, preview: bool = True
+    ) -> PaginationResult:
+        """
+        List all conversation sessions for an agent created by the calling user.
+
+        Note: Only lists sessions created by this client, not sessions created
+        in AIP Agent Studio.
+
+        Args:
+            agent_rid: Agent Resource Identifier
+            config: Pagination configuration
+            preview: Enable preview mode (default: True)
+
+        Returns:
+            PaginationResult with sessions and metadata
+
+        Example:
+            >>> config = PaginationConfig(page_size=20, max_pages=2)
+            >>> result = service.list_sessions(agent_rid, config)
+            >>> print(f"Found {result.metadata.items_fetched} sessions")
+        """
+        try:
+            iterator = self.service.Agent.Session.list(
+                agent_rid, page_size=config.page_size, preview=preview
+            )
+            return self._paginate_iterator(iterator, config)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list sessions for agent {agent_rid}: {e}")
+
+    def get_session(
+        self, agent_rid: str, session_rid: str, preview: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Get details of a specific conversation session.
+
+        Args:
+            agent_rid: Agent Resource Identifier
+            session_rid: Session Resource Identifier
+            preview: Enable preview mode (default: True)
+
+        Returns:
+            Session information dictionary containing:
+            - rid: Session resource identifier
+            - agent_rid: Associated agent RID
+            - agent_version: Agent version used in session
+            - metadata: Session metadata (e.g., title)
+
+        Raises:
+            RuntimeError: If the operation fails
+        """
+        try:
+            session = self.service.Agent.Session.get(
+                agent_rid, session_rid, preview=preview
+            )
+            return self._serialize_response(session)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get session {session_rid} for agent {agent_rid}: {e}"
+            )
+
+    # ===== Version Operations =====
+
+    def list_versions(
+        self, agent_rid: str, config: PaginationConfig, preview: bool = True
+    ) -> PaginationResult:
+        """
+        List all versions for an AIP Agent.
+
+        Versions are returned in descending order (most recent first).
+
+        Args:
+            agent_rid: Agent Resource Identifier
+            config: Pagination configuration
+            preview: Enable preview mode (default: True)
+
+        Returns:
+            PaginationResult with agent versions and metadata
+
+        Example:
+            >>> config = PaginationConfig(page_size=10, fetch_all=True)
+            >>> result = service.list_versions(agent_rid, config)
+            >>> for version in result.data:
+            ...     print(f"Version {version['string']}")
+        """
+        try:
+            iterator = self.service.Agent.AgentVersion.list(
+                agent_rid, page_size=config.page_size, preview=preview
+            )
+            return self._paginate_iterator(iterator, config)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list versions for agent {agent_rid}: {e}")

pltr/services/base.py

@@ -2,13 +2,20 @@
 Base service class for Foundry API wrappers.
 """
 
-from typing import Any, Optional, Dict
+from typing import Any, Optional, Dict, Callable, Iterator
 from abc import ABC, abstractmethod
+import json
 import requests
 
 from ..auth.manager import AuthManager
 from ..auth.storage import CredentialStorage
 from ..config.profiles import ProfileManager
+from ..utils.pagination import (
+    PaginationConfig,
+    PaginationResult,
+    IteratorPaginationHandler,
+    ResponsePaginationHandler,
+)
 
 
 class BaseService(ABC):
@@ -127,3 +134,99 @@ class BaseService(ABC):
         response.raise_for_status()
 
         return response
+
+    def _paginate_iterator(
+        self,
+        iterator: Iterator[Any],
+        config: PaginationConfig,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> PaginationResult:
+        """
+        Handle pagination for iterator-based SDK methods.
+
+        Args:
+            iterator: Iterator returned from SDK (e.g., ResourceIterator)
+            config: Pagination configuration
+            progress_callback: Optional progress callback
+
+        Returns:
+            PaginationResult with collected items and metadata
+
+        Example:
+            >>> iterator = self.service.Dataset.File.list(dataset_rid, page_size=20)
+            >>> result = self._paginate_iterator(iterator, config)
+        """
+        handler = IteratorPaginationHandler()
+        return handler.collect_pages(iterator, config, progress_callback)
+
+    def _paginate_response(
+        self,
+        fetch_fn: Callable[[Optional[str]], Dict[str, Any]],
+        config: PaginationConfig,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> PaginationResult:
+        """
+        Handle pagination for response-based SDK methods.
+
+        Args:
+            fetch_fn: Function that accepts page_token and returns dict with
+                     'data' and 'next_page_token' keys
+            config: Pagination configuration
+            progress_callback: Optional progress callback
+
+        Returns:
+            PaginationResult with collected items and metadata
+
+        Example:
+            >>> def fetch(token):
+            ...     response = self.service.User.list(page_token=token)
+            ...     return self._serialize_response(response)
+            >>> result = self._paginate_response(fetch, config)
+        """
+        handler = ResponsePaginationHandler()
+        return handler.collect_pages(fetch_fn, config, progress_callback)
+
+    def _serialize_response(self, response: Any) -> Dict[str, Any]:
+        """
+        Convert response object to serializable dictionary.
+
+        This handles various SDK response types including Pydantic models,
+        regular objects, and primitives.
+
+        Args:
+            response: Response object from SDK
+
+        Returns:
+            Serializable dictionary representation
+
+        Note:
+            This method was moved from AdminService to provide consistent
+            serialization across all services.
+        """
+        if response is None:
+            return {}
+
+        # Handle different response types
+        if hasattr(response, "dict"):
+            # Pydantic models
+            return response.dict()
+        elif hasattr(response, "__dict__"):
+            # Regular objects
+            result = {}
+            for key, value in response.__dict__.items():
+                if not key.startswith("_"):
+                    try:
+                        # Try to serialize the value
+                        json.dumps(value)
+                        result[key] = value
+                    except (TypeError, ValueError):
+                        # Convert non-serializable values to string
+                        result[key] = str(value)
+            return result
+        else:
+            # Primitive types or already serializable
+            try:
+                json.dumps(response)
+                return response
+            except (TypeError, ValueError):
+                return {"data": str(response)}