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

examples/sdk_usage.py

@@ -0,0 +1,158 @@
+"""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)
+    orgs = []
+
+    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:")
+        orgs = list(orgs_client.list_all())
+        for org in orgs:
+            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, CodegenError, NotFoundError
+
+    # 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()

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
   "codeowners>=0.6.0",
   "unidiff>=0.7.5",
   "datamodel-code-generator>=0.26.5",
-  "mcp[cli]",
+  "mcp[cli]==1.9.4",
+  "fastmcp>=2.9.0",
   # Utility dependencies
   "colorlog>=6.9.0",
   "psutil>=5.8.0",

src/codegen/agents/__init__.py

@@ -1,5 +1,10 @@
 """Codegen Agent API module."""
 
 from codegen.agents.agent import Agent
+from codegen.organizations import Organizations
 
-__all__ = ["Agent"]
+# Import SDK components for backward compatibility and convenience
+from codegen.sdk import CodegenSDK
+from codegen.users import Users
+
+__all__ = ["Agent", "CodegenSDK", "Organizations", "Users"]

src/codegen/agents/agent.py

@@ -6,8 +6,10 @@
 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,66 @@
         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}"
-        )
+        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:
+            status_code = getattr(e, "status", None)
+            if not isinstance(status_code, int):
+                status_code = 500  # Default to server error if status is not available
+            error = handle_api_error(status_code, 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"]
 
-        # 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
+    def is_running(self) -> bool:
+        """Check if the agent task is currently running."""
+        return self.status in ["running", "pending"]
 
-        self.status = job_dict.get("status")
-        self.result = job_dict.get("result")
+    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 | None, org_id: int | None = None, base_url: str | None = CODEGEN_BASE_API_URL):
+    def __init__(self, token: str | None, org_id: int | str | None = None, base_url: str | None = 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
@@ -70,26 +104,60 @@
             prompt: The instruction for the agent to execute
 
         Returns:
-            Job: A job object representing the agent run
-        """
-        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
+            AgentTask: A task object representing the agent run
 
-        job = AgentTask(agent_run_response, self.api_client, self.org_id)
-        self.current_job = job
-        return job
+        Raises:
+            CodegenError: If the API request fails
+        """
+        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:
+            status_code = getattr(e, "status", None)
+            if not isinstance(status_code, int):
+                status_code = 500  # Default to server error if status is not available
+            error = handle_api_error(status_code, str(e), getattr(e, "body", None))
+            raise error from e
 
     def get_status(self) -> dict[str, Any] | None:
         """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: 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:
+            status_code = getattr(e, "status", None)
+            if not isinstance(status_code, int):
+                status_code = 500  # Default to server error if status is not available
+            error = handle_api_error(status_code, str(e), getattr(e, "body", None))
+            raise error from e

src/codegen/cli/api/client.py

@@ -5,8 +5,6 @@
 from pydantic import BaseModel
 from rich import print as rprint
 
-from codegen.cli.api.endpoints import IDENTIFY_ENDPOINT
-from codegen.cli.api.schemas import IdentifyResponse
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import InvalidTokenError, ServerError
 
@@ -16,11 +14,13 @@
 
 class AuthContext(BaseModel):
     """Authentication context model."""
+
     status: str
 
 
 class Identity(BaseModel):
     """User identity model."""
+
     auth_context: AuthContext
 
 

src/codegen/cli/auth/decorators.py

@@ -39,7 +39,7 @@ def wrapper(*args, **kwargs):
 
     # Remove the session parameter from the wrapper's signature so Typer doesn't see it
     sig = inspect.signature(f)
-    new_params = [param for name, param in sig.parameters.items() if name != 'session']
+    new_params = [param for name, param in sig.parameters.items() if name != "session"]
     new_sig = sig.replace(parameters=new_params)
     wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
 

src/codegen/cli/auth/session.ipynb

@@ -11,7 +11,7 @@
     "\n",
     "\n",
     "# Create a session with the current directory as repo_path\n",
-    "session = CodegenSession(repo_path=Path('.'))\n",
+    "session = CodegenSession(repo_path=Path(\".\"))\n",
     "print(f\"Session: {session}\")\n",
     "print(f\"Repo path: {session.repo_path}\")\n",
     "print(f\"Config: {session.config}\")\n",