chore(typing): Add Pydantic models for all tools

Author: timothedelion

packages/gg_api_core/src/gg_api_core/tools/find_current_source_id.py

@@ -1,13 +1,41 @@
 from typing import Any
 import logging
 import subprocess
+from pydantic import BaseModel, Field
 from gg_api_core.utils import get_client, parse_repo_url
 
 logger = logging.getLogger(__name__)
 
 
+class SourceCandidate(BaseModel):
+    """A candidate source that might match the repository."""
+    id: str = Field(description="Source ID")
+    url: str | None = Field(default=None, description="Repository URL")
+    name: str | None = Field(default=None, description="Repository name")
+    monitored: bool | None = Field(default=None, description="Whether source is monitored")
+    deleted_at: str | None = Field(default=None, description="Deletion timestamp if deleted")
 
-async def find_current_source_id() -> dict[str, Any]:
+
+class FindCurrentSourceIdResult(BaseModel):
+    """Successful result from finding source ID."""
+    repository_name: str = Field(description="Detected repository name")
+    source_id: str | None = Field(default=None, description="GitGuardian source ID (if exact match)")
+    source: dict[str, Any] | None = Field(default=None, description="Full source information (if exact match)")
+    message: str | None = Field(default=None, description="Status or informational message")
+    suggestion: str | None = Field(default=None, description="Suggestions for next steps")
+    candidates: list[SourceCandidate] | None = Field(default=None, description="List of candidate sources (if no exact match)")
+
+
+class FindCurrentSourceIdError(BaseModel):
+    """Error result from finding source ID."""
+    error: str = Field(description="Error message")
+    repository_name: str | None = Field(default=None, description="Repository name if detected")
+    details: str | None = Field(default=None, description="Additional error details")
+    message: str | None = Field(default=None, description="User-friendly message")
+    suggestion: str | None = Field(default=None, description="Suggestions for resolving the error")
+
+
+async def find_current_source_id() -> FindCurrentSourceIdResult | FindCurrentSourceIdError:
     """
     Find the GitGuardian source_id for the current repository.
 
@@ -19,12 +47,20 @@ async def find_current_source_id() -> dict[str, Any]:
     5. If no exact match, returns all search results for the model to choose from
 
     Returns:
-        A dictionary containing:
-        - repository_name: The detected repository name
-        - source_id: The GitGuardian source ID (if exact match found)
-        - source: Full source information from GitGuardian (if exact match found)
-        - candidates: List of candidate sources (if no exact match but potential matches found)
-        - error: Error message if something went wrong
+        FindCurrentSourceIdResult: Pydantic model containing:
+            - repository_name: The detected repository name
+            - source_id: The GitGuardian source ID (if exact match found)
+            - source: Full source information from GitGuardian (if exact match found)
+            - message: Status or informational message
+            - suggestion: Suggestions for next steps
+            - candidates: List of SourceCandidate objects (if no exact match but potential matches found)
+
+        FindCurrentSourceIdError: Pydantic model containing:
+            - error: Error message
+            - repository_name: Repository name if detected
+            - details: Additional error details
+            - message: User-friendly message
+            - suggestion: Suggestions for resolving the error
     """
     client = get_client()
     logger.debug("Finding source_id for current repository")
@@ -42,21 +78,21 @@ async def find_current_source_id() -> dict[str, Any]:
             remote_url = result.stdout.strip()
             logger.debug(f"Found remote URL: {remote_url}")
         except subprocess.CalledProcessError as e:
-            return {
-                "error": "Not a git repository or no remote 'origin' configured",
-                "details": str(e),
-            }
+            return FindCurrentSourceIdError(
+                error="Not a git repository or no remote 'origin' configured",
+                details=str(e),
+            )
         except subprocess.TimeoutExpired:
-            return {"error": "Git command timed out"}
+            return FindCurrentSourceIdError(error="Git command timed out")
 
         # Parse repository name from remote URL
         repository_name = parse_repo_url(remote_url)
 
         if not repository_name:
-            return {
-                "error": f"Could not parse repository URL: {remote_url}",
-                "details": "The URL format is not recognized. Supported platforms: GitHub, GitLab (Cloud & Self-hosted), Bitbucket (Cloud & Data Center), Azure DevOps",
-            }
+            return FindCurrentSourceIdError(
+                error=f"Could not parse repository URL: {remote_url}",
+                details="The URL format is not recognized. Supported platforms: GitHub, GitLab (Cloud & Self-hosted), Bitbucket (Cloud & Data Center), Azure DevOps",
+            )
 
         logger.info(f"Detected repository name: {repository_name}")
 
@@ -67,31 +103,31 @@ async def find_current_source_id() -> dict[str, Any]:
         if isinstance(result, dict):
             source_id = result.get("id")
             logger.info(f"Found exact match with source_id: {source_id}")
-            return {
-                "repository_name": repository_name,
-                "source_id": source_id,
-                "source": result,
-                "message": f"Successfully found exact match for GitGuardian source: {repository_name}",
-            }
+            return FindCurrentSourceIdResult(
+                repository_name=repository_name,
+                source_id=source_id,
+                source=result,
+                message=f"Successfully found exact match for GitGuardian source: {repository_name}",
+            )
 
         # Handle multiple candidates (list result)
         elif isinstance(result, list) and len(result) > 0:
             logger.info(f"Found {len(result)} candidate sources for repository: {repository_name}")
-            return {
-                "repository_name": repository_name,
-                "message": f"No exact match found for '{repository_name}', but found {len(result)} potential matches.",
-                "suggestion": "Review the candidates below and determine which source best matches the current repository based on the name and URL.",
-                "candidates": [
-                    {
-                        "id": source.get("id"),
-                        "url": source.get("url"),
-                        "name": source.get("full_name") or source.get("name"),
-                        "monitored": source.get("monitored"),
-                        "deleted_at": source.get("deleted_at"),
-                    }
+            return FindCurrentSourceIdResult(
+                repository_name=repository_name,
+                message=f"No exact match found for '{repository_name}', but found {len(result)} potential matches.",
+                suggestion="Review the candidates below and determine which source best matches the current repository based on the name and URL.",
+                candidates=[
+                    SourceCandidate(
+                        id=source.get("id"),
+                        url=source.get("url"),
+                        name=source.get("full_name") or source.get("name"),
+                        monitored=source.get("monitored"),
+                        deleted_at=source.get("deleted_at"),
+                    )
                     for source in result
                 ],
-            }
+            )
 
         # No matches found at all
         else:
@@ -105,39 +141,39 @@ async def find_current_source_id() -> dict[str, Any]:
                 if isinstance(fallback_result, dict):
                     source_id = fallback_result.get("id")
                     logger.info(f"Found match using repo name only, source_id: {source_id}")
-                    return {
-                        "repository_name": repository_name,
-                        "source_id": source_id,
-                        "source": fallback_result,
-                        "message": f"Found match using repository name '{repo_only}' (without organization prefix)",
-                    }
+                    return FindCurrentSourceIdResult(
+                        repository_name=repository_name,
+                        source_id=source_id,
+                        source=fallback_result,
+                        message=f"Found match using repository name '{repo_only}' (without organization prefix)",
+                    )
                 elif isinstance(fallback_result, list) and len(fallback_result) > 0:
                     logger.info(f"Found {len(fallback_result)} candidates using repo name only")
-                    return {
-                        "repository_name": repository_name,
-                        "message": f"No exact match for '{repository_name}', but found {len(fallback_result)} potential matches using repo name '{repo_only}'.",
-                        "suggestion": "Review the candidates below and determine which source best matches the current repository.",
-                        "candidates": [
-                            {
-                                "id": source.get("id"),
-                                "url": source.get("url"),
-                                "name": source.get("full_name") or source.get("name"),
-                                "monitored": source.get("monitored"),
-                                "deleted_at": source.get("deleted_at"),
-                            }
+                    return FindCurrentSourceIdResult(
+                        repository_name=repository_name,
+                        message=f"No exact match for '{repository_name}', but found {len(fallback_result)} potential matches using repo name '{repo_only}'.",
+                        suggestion="Review the candidates below and determine which source best matches the current repository.",
+                        candidates=[
+                            SourceCandidate(
+                                id=source.get("id"),
+                                url=source.get("url"),
+                                name=source.get("full_name") or source.get("name"),
+                                monitored=source.get("monitored"),
+                                deleted_at=source.get("deleted_at"),
+                            )
                             for source in fallback_result
                         ],
-                    }
+                    )
 
             # Absolutely no matches found
             logger.warning(f"No sources found for repository: {repository_name}")
-            return {
-                "repository_name": repository_name,
-                "error": f"Repository '{repository_name}' not found in GitGuardian",
-                "message": "The repository may not be connected to GitGuardian, or you may not have access to it.",
-                "suggestion": "Check that the repository is properly connected to GitGuardian and that your account has access to it.",
-            }
+            return FindCurrentSourceIdError(
+                repository_name=repository_name,
+                error=f"Repository '{repository_name}' not found in GitGuardian",
+                message="The repository may not be connected to GitGuardian, or you may not have access to it.",
+                suggestion="Check that the repository is properly connected to GitGuardian and that your account has access to it.",
+            )
 
     except Exception as e:
         logger.error(f"Error finding source_id: {str(e)}")
-        return {"error": f"Failed to find source_id: {str(e)}"}
+        return FindCurrentSourceIdError(error=f"Failed to find source_id: {str(e)}")

packages/gg_api_core/src/gg_api_core/tools/generate_honey_token.py

@@ -21,15 +21,42 @@ class GenerateHoneytokenParams(BaseModel):
     )
 
 
-async def generate_honeytoken(params: GenerateHoneytokenParams) -> dict[str, Any]:
+class GenerateHoneytokenResult(BaseModel):
+    """Result from generating or retrieving a honeytoken."""
+    model_config = {"extra": "allow"}  # Allow additional fields from API
+
+    id: str = Field(description="Honeytoken ID")
+    name: str | None = Field(default=None, description="Honeytoken name")
+    description: str | None = Field(default=None, description="Honeytoken description")
+    type: str | None = Field(default=None, description="Honeytoken type")
+    status: str | None = Field(default=None, description="Honeytoken status")
+    token: dict[str, Any] | None = Field(default=None, description="Token details if show_token=True")
+    injection_recommendations: dict[str, Any] | None = Field(default=None, description="Injection recommendations")
+
+
+async def generate_honeytoken(params: GenerateHoneytokenParams) -> GenerateHoneytokenResult:
     """
     Generate an AWS GitGuardian honeytoken and get injection recommendations.
 
+    If new_token is False, attempts to retrieve an existing active honeytoken created by the current user
+    instead of generating a new one. If no existing token is found, a new one will be created.
+
     Args:
         params: GenerateHoneytokenParams model containing honeytoken configuration
 
     Returns:
-        Honeytoken data and injection recommendations
+        GenerateHoneytokenResult: Pydantic model containing:
+            - id: Honeytoken ID
+            - name: Honeytoken name
+            - description: Honeytoken description
+            - type: Honeytoken type
+            - status: Honeytoken status
+            - token: Token details (if show_token=True was used)
+            - injection_recommendations: Instructions for injecting the honeytoken
+            - Additional fields from the API response
+
+    Raises:
+        ToolError: If the honeytoken generation or retrieval fails
     """
     client = get_client()
     logger.debug(f"Processing honeytoken request with name: {params.name}, new_token: {params.new_token}")
@@ -67,7 +94,7 @@ async def generate_honeytoken(params: GenerateHoneytokenParams) -> dict[str, Any
                     if honeytoken_id:
                         detailed_token = await client.get_honeytoken(honeytoken_id, show_token=True)
                         logger.debug(f"Retrieved existing honeytoken with ID: {honeytoken_id}")
-                        return detailed_token
+                        return GenerateHoneytokenResult(**detailed_token)
 
                 logger.debug("No suitable existing honeytokens found, creating a new one")
             else:
@@ -95,7 +122,7 @@ async def generate_honeytoken(params: GenerateHoneytokenParams) -> dict[str, Any
             "instructions": "Add the honeytoken to your codebase in configuration files, environment variables, or code comments to detect unauthorized access."
         }
 
-        return result
+        return GenerateHoneytokenResult(**result)
     except Exception as e:
         logger.error(f"Error generating honeytoken: {str(e)}")
         raise ToolError(f"Failed to generate honeytoken: {str(e)}")

packages/gg_api_core/src/gg_api_core/tools/list_honey_tokens.py

@@ -25,15 +25,26 @@ class ListHoneytokensParams(BaseModel):
     mine: bool = Field(default=False, description="If True, fetch honeytokens created by the current user")
 
 
-async def list_honeytokens(params: ListHoneytokensParams) -> list[dict[str, Any]]:
+class ListHoneytokensResult(BaseModel):
+    """Result from listing honeytokens."""
+    honeytokens: list[dict[str, Any]] = Field(description="List of honeytoken objects")
+
+
+async def list_honeytokens(params: ListHoneytokensParams) -> ListHoneytokensResult:
     """
     List honeytokens from the GitGuardian dashboard with filtering options.
 
+    If mine=True, filters honeytokens to show only those created by the current user.
+
     Args:
         params: ListHoneytokensParams model containing all filtering options
 
     Returns:
-        List of honeytokens matching the specified criteria
+        ListHoneytokensResult: Pydantic model containing:
+            - honeytokens: List of honeytoken objects matching the specified criteria
+
+    Raises:
+        ToolError: If the listing operation fails
     """
     client = get_client()
     logger.debug("Listing honeytokens with filters")
@@ -79,7 +90,7 @@ async def list_honeytokens(params: ListHoneytokensParams) -> list[dict[str, Any]
             honeytokens = result
 
         logger.debug(f"Found {len(honeytokens)} honeytokens")
-        return honeytokens
+        return ListHoneytokensResult(honeytokens=honeytokens)
     except Exception as e:
         logger.error(f"Error listing honeytokens: {str(e)}")
         raise ToolError(str(e))

packages/gg_api_core/src/gg_api_core/tools/list_repo_incidents.py

@@ -151,10 +151,19 @@ async def list_repo_incidents(params: ListRepoIncidentsParams) -> ListRepoIncide
     By default, incidents tagged with TEST_FILE or FALSE_POSITIVE are excluded. Pass exclude_tags=[] to disable this filtering.
 
     Args:
-        params: ListRepoIncidentsParams model containing all filtering options
+        params: ListRepoIncidentsParams model containing all filtering options.
+               Either repository_name or source_id must be provided.
 
     Returns:
-        List of incidents and occurrences matching the specified criteria
+        ListRepoIncidentsResult: Pydantic model containing:
+            - source_id: Source ID of the repository
+            - incidents: List of incident objects
+            - total_count: Total number of incidents
+            - next_cursor: Pagination cursor (if applicable)
+            - applied_filters: Dictionary of filters that were applied
+            - suggestion: Suggestions for interpreting or modifying results
+
+        ListRepoIncidentsError: Pydantic model with error message if the operation fails
     """
     client = get_client()