feat: Enhanced SDK support with Organizations, Users, and unified CodegenSDK

- Added comprehensive exception handling with custom error types
- Implemented Organizations API client with pagination support
- Implemented Users API client with search and filtering capabilities
- Enhanced Agent class with better error handling and helper methods
- Created unified CodegenSDK class as main entry point
- Added comprehensive examples and documentation
- Maintained backward compatibility with existing Agent interface

This addresses CG-18852 by providing a more complete set of operations
with intuitive objects that make sense, building on the existing
OpenAPI spec from codegen-api-client.

Author: codegen-sh[bot]

examples/sdk_usage.py

@@ -0,0 +1,155 @@
+"""
+Example usage of the enhanced Codegen SDK.
+
+This example demonstrates how to use the new SDK components:
+- CodegenSDK: Unified entry point for all API functionality
+- Organizations: Managing organizations
+- Users: Managing users within organizations
+- Enhanced Agent: Improved agent functionality with better error handling
+"""
+
+import os
+from codegen.agents import Agent, CodegenSDK, Organizations, Users
+
+# Example 1: Using the unified SDK
+def example_unified_sdk():
+    """Example using the unified CodegenSDK."""
+    print("=== Unified SDK Example ===")
+    
+    # Initialize the SDK with your API token
+    sdk = CodegenSDK(token=os.getenv("CODEGEN_API_TOKEN"))
+    
+    # Use agents
+    print("Running an agent task...")
+    task = sdk.agents.run("Analyze the codebase and suggest improvements")
+    print(f"Task created: {task.id}")
+    print(f"Status: {task.status}")
+    print(f"Web URL: {task.web_url}")
+    
+    # List organizations
+    print("\nListing organizations...")
+    orgs = sdk.organizations.list()
+    for org in orgs:
+        print(f"- {org.name} (ID: {org.id})")
+    
+    # Get users in the first organization
+    if orgs:
+        print(f"\nListing users in organization '{orgs[0].name}'...")
+        users = sdk.users(org_id=orgs[0].id).list()
+        for user in users:
+            print(f"- {user.github_username} ({user.email})")
+
+
+# Example 2: Using individual SDK components
+def example_individual_components():
+    """Example using individual SDK components."""
+    print("\n=== Individual Components Example ===")
+    
+    token = os.getenv("CODEGEN_API_TOKEN")
+    
+    # Enhanced Agent with better error handling
+    print("Using enhanced Agent...")
+    agent = Agent(token=token)
+    
+    try:
+        task = agent.run("Create a simple Python function")
+        print(f"Agent task: {task.id}")
+        
+        # Check task status with new helper methods
+        if task.is_running():
+            print("Task is currently running...")
+        elif task.is_completed():
+            print("Task completed!")
+            if task.is_successful():
+                print("Task was successful!")
+            elif task.is_failed():
+                print("Task failed.")
+        
+        # Get task by ID
+        retrieved_task = agent.get_task(task.id)
+        print(f"Retrieved task: {retrieved_task.id}")
+        
+    except Exception as e:
+        print(f"Error: {e}")
+    
+    # Organizations management
+    print("\nUsing Organizations...")
+    orgs_client = Organizations(token=token)
+    
+    try:
+        # List with pagination
+        page_result = orgs_client.get_page(page=1, page_size=10)
+        print(f"Organizations page: {page_result['page']}/{page_result['pages']}")
+        print(f"Total organizations: {page_result['total']}")
+        
+        # Iterate through all organizations
+        print("All organizations:")
+        for org in orgs_client.list_all():
+            print(f"- {org.name} (ID: {org.id})")
+            print(f"  Settings: {org.settings}")
+            
+    except Exception as e:
+        print(f"Error: {e}")
+    
+    # Users management
+    print("\nUsing Users...")
+    if orgs:
+        users_client = Users(token=token, org_id=orgs[0].id)
+        
+        try:
+            # List users with pagination
+            page_result = users_client.get_page(page=1, page_size=5)
+            print(f"Users page: {page_result['page']}/{page_result['pages']}")
+            print(f"Total users: {page_result['total']}")
+            
+            # Find specific users
+            user = users_client.find_by_github_username("example-user")
+            if user:
+                print(f"Found user: {user.github_username}")
+            
+            # Get user by ID
+            if page_result['items']:
+                first_user = page_result['items'][0]
+                retrieved_user = users_client.get(first_user.id)
+                print(f"Retrieved user: {retrieved_user.github_username}")
+                
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+# Example 3: Error handling
+def example_error_handling():
+    """Example demonstrating error handling."""
+    print("\n=== Error Handling Example ===")
+    
+    from codegen.exceptions import AuthenticationError, NotFoundError, CodegenError
+    
+    # Invalid token example
+    try:
+        sdk = CodegenSDK(token="invalid-token")
+        orgs = sdk.organizations.list()
+    except AuthenticationError as e:
+        print(f"Authentication failed: {e.message}")
+        print(f"Status code: {e.status_code}")
+    except CodegenError as e:
+        print(f"API error: {e.message}")
+    
+    # Not found example
+    try:
+        sdk = CodegenSDK(token=os.getenv("CODEGEN_API_TOKEN"))
+        user = sdk.users().get(user_id=99999)  # Non-existent user
+    except NotFoundError as e:
+        print(f"User not found: {e.message}")
+    except CodegenError as e:
+        print(f"API error: {e.message}")
+
+
+if __name__ == "__main__":
+    # Make sure to set your API token
+    if not os.getenv("CODEGEN_API_TOKEN"):
+        print("Please set the CODEGEN_API_TOKEN environment variable")
+        exit(1)
+    
+    example_unified_sdk()
+    example_individual_components()
+    example_error_handling()

src/codegen/agents/__init__.py

@@ -2,4 +2,9 @@
 
 from codegen.agents.agent import Agent
 
-__all__ = ["Agent"]
+# Import SDK components for backward compatibility and convenience
+from codegen.sdk import CodegenSDK
+from codegen.organizations import Organizations
+from codegen.users import Users
+
+__all__ = ["Agent", "CodegenSDK", "Organizations", "Users"]

src/codegen/agents/agent.py

@@ -1,13 +1,15 @@
 import os
-from typing import Any
+from typing import Any, Dict, Optional, Union
 
 from codegen_api_client.api.agents_api import AgentsApi
 from codegen_api_client.api_client import ApiClient
 from codegen_api_client.configuration import Configuration
 from codegen_api_client.models.agent_run_response import AgentRunResponse
 from codegen_api_client.models.create_agent_run_input import CreateAgentRunInput
+from codegen_api_client.rest import ApiException
 
 from codegen.agents.constants import CODEGEN_BASE_API_URL
+from codegen.exceptions import handle_api_error
 
 
 class AgentTask:
@@ -23,34 +25,69 @@ def __init__(self, task_data: AgentRunResponse, api_client: ApiClient, org_id: i
         self._agents_api = AgentsApi(api_client)
 
     def refresh(self) -> None:
-        """Refresh the job status from the API."""
+        """Refresh the job status from the API.
+        
+        Raises:
+            CodegenError: If the API request fails
+        """
         if self.id is None:
             return
 
-        job_data = self._agents_api.get_agent_run_v1_organizations_org_id_agent_run_agent_run_id_get(
-            agent_run_id=int(self.id), org_id=int(self.org_id), authorization=f"Bearer {self._api_client.configuration.access_token}"
-        )
-
-        # Convert API response to dict for attribute access
-        job_dict = {}
-        if hasattr(job_data, "__dict__"):
-            job_dict = job_data.__dict__
-        elif isinstance(job_data, dict):
-            job_dict = job_data
-
-        self.status = job_dict.get("status")
-        self.result = job_dict.get("result")
+        try:
+            job_data = self._agents_api.get_agent_run_v1_organizations_org_id_agent_run_agent_run_id_get(
+                agent_run_id=int(self.id), org_id=int(self.org_id), authorization=f"Bearer {self._api_client.configuration.access_token}"
+            )
+
+            # Convert API response to dict for attribute access
+            job_dict = {}
+            if hasattr(job_data, "__dict__"):
+                job_dict = job_data.__dict__
+            elif isinstance(job_data, dict):
+                job_dict = job_data
+
+            self.status = job_dict.get("status")
+            self.result = job_dict.get("result")
+        except ApiException as e:
+            error = handle_api_error(e.status, str(e), getattr(e, 'body', None))
+            raise error from e
+    
+    def is_completed(self) -> bool:
+        """Check if the agent task is completed."""
+        return self.status in ['completed', 'failed', 'cancelled']
+    
+    def is_running(self) -> bool:
+        """Check if the agent task is currently running."""
+        return self.status in ['running', 'pending']
+    
+    def is_successful(self) -> bool:
+        """Check if the agent task completed successfully."""
+        return self.status == 'completed'
+    
+    def is_failed(self) -> bool:
+        """Check if the agent task failed."""
+        return self.status == 'failed'
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert agent task to dictionary."""
+        return {
+            'id': self.id,
+            'org_id': self.org_id,
+            'status': self.status,
+            'result': self.result,
+            'web_url': self.web_url
+        }
 
 
 class Agent:
     """API client for interacting with Codegen AI agents."""
 
-    def __init__(self, token: str, org_id: int | None = None, base_url: str | None = CODEGEN_BASE_API_URL):
+    def __init__(self, token: str, org_id: Optional[Union[int, str]] = None, base_url: str = CODEGEN_BASE_API_URL):
         """Initialize a new Agent client.
 
         Args:
             token: API authentication token
             org_id: Optional organization ID. If not provided, default org will be used.
+            base_url: Base URL for the API (defaults to production)
         """
         self.token = token
         self.org_id = org_id or int(os.environ.get("CODEGEN_ORG_ID", "1"))  # Default to org ID 1 if not specified
@@ -61,7 +98,7 @@ def __init__(self, token: str, org_id: int | None = None, base_url: str | None =
         self.agents_api = AgentsApi(self.api_client)
 
         # Current job
-        self.current_job: AgentTask | None = None
+        self.current_job: Optional[AgentTask] = None
 
     def run(self, prompt: str) -> AgentTask:
         """Run an agent with the given prompt.
@@ -70,26 +107,58 @@ def run(self, prompt: str) -> AgentTask:
             prompt: The instruction for the agent to execute
 
         Returns:
-            Job: A job object representing the agent run
+            AgentTask: A task object representing the agent run
+            
+        Raises:
+            CodegenError: If the API request fails
         """
-        run_input = CreateAgentRunInput(prompt=prompt)
-        agent_run_response = self.agents_api.create_agent_run_v1_organizations_org_id_agent_run_post(
-            org_id=int(self.org_id), create_agent_run_input=run_input, authorization=f"Bearer {self.token}", _headers={"Content-Type": "application/json"}
-        )
-        # Convert API response to dict for Job initialization
-
-        job = AgentTask(agent_run_response, self.api_client, self.org_id)
-        self.current_job = job
-        return job
-
-    def get_status(self) -> dict[str, Any] | None:
+        try:
+            run_input = CreateAgentRunInput(prompt=prompt)
+            agent_run_response = self.agents_api.create_agent_run_v1_organizations_org_id_agent_run_post(
+                org_id=int(self.org_id), create_agent_run_input=run_input, authorization=f"Bearer {self.token}", _headers={"Content-Type": "application/json"}
+            )
+
+            job = AgentTask(agent_run_response, self.api_client, self.org_id)
+            self.current_job = job
+            return job
+        except ApiException as e:
+            error = handle_api_error(e.status, str(e), getattr(e, 'body', None))
+            raise error from e
+
+    def get_status(self) -> Optional[Dict[str, Any]]:
         """Get the status of the current job.
 
         Returns:
             dict: A dictionary containing job status information,
                   or None if no job has been run.
+                  
+        Raises:
+            CodegenError: If the API request fails
         """
         if self.current_job:
             self.current_job.refresh()
             return {"id": self.current_job.id, "status": self.current_job.status, "result": self.current_job.result, "web_url": self.current_job.web_url}
         return None
+    
+    def get_task(self, task_id: Union[int, str]) -> AgentTask:
+        """Get a specific agent task by ID.
+        
+        Args:
+            task_id: Agent task ID to retrieve
+            
+        Returns:
+            AgentTask: The requested agent task
+            
+        Raises:
+            CodegenError: If the API request fails or task is not found
+        """
+        try:
+            response = self.agents_api.get_agent_run_v1_organizations_org_id_agent_run_agent_run_id_get(
+                org_id=int(self.org_id),
+                agent_run_id=int(task_id),
+                authorization=f"Bearer {self.token}"
+            )
+            return AgentTask(response, self.api_client, self.org_id)
+        except ApiException as e:
+            error = handle_api_error(e.status, str(e), getattr(e, 'body', None))
+            raise error from e