chore(typing): Add Pydantic models for tools return : remediate and list

Author: timothedelion

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

@@ -128,7 +128,22 @@ class ListRepoIncidentsParams(BaseModel):
     validity: list[str] | None = Field(default=DEFAULT_VALIDITIES, description="Filter by validity (list of validity names)")
 
 
-async def list_repo_incidents(params: ListRepoIncidentsParams) -> dict[str, Any]:
+class ListRepoIncidentsResult(BaseModel):
+    """Result from listing repository incidents."""
+    source_id: str | None = Field(default=None, description="Source ID of the repository")
+    incidents: list[dict[str, Any]] = Field(default_factory=list, description="List of incident objects")
+    total_count: int = Field(description="Total number of incidents")
+    next_cursor: str | None = Field(default=None, description="Pagination cursor for next page")
+    applied_filters: dict[str, Any] = Field(default_factory=dict, description="Filters that were applied to the query")
+    suggestion: str = Field(default="", description="Suggestions for interpreting or modifying the results")
+
+
+class ListRepoIncidentsError(BaseModel):
+    """Error result from listing repository incidents."""
+    error: str = Field(description="Error message")
+
+
+async def list_repo_incidents(params: ListRepoIncidentsParams) -> ListRepoIncidentsResult | ListRepoIncidentsError:
     """
     List secret incidents or occurrences related to a specific repository.
 
@@ -145,7 +160,7 @@ async def list_repo_incidents(params: ListRepoIncidentsParams) -> dict[str, Any]
 
     # Validate that at least one of repository_name or source_id is provided
     if not params.repository_name and not params.source_id:
-        return {"error": "Either repository_name or source_id must be provided"}
+        return ListRepoIncidentsError(error="Either repository_name or source_id must be provided")
 
     logger.debug(f"Listing incidents with repository_name={params.repository_name}, source_id={params.source_id}")
 
@@ -185,64 +200,54 @@ async def list_repo_incidents(params: ListRepoIncidentsParams) -> dict[str, Any]
                 incidents_result = await client.paginate_all(f"/sources/{params.source_id}/incidents/secrets", api_params)
                 if isinstance(incidents_result, list):
                     count = len(incidents_result)
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": incidents_result,
-                        "total_count": count,
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, count),
-                    }
+                    return ListRepoIncidentsResult(
+                        source_id=params.source_id,
+                        incidents=incidents_result,
+                        total_count=count,
+                        applied_filters=_build_filter_info(params),
+                        suggestion=_build_suggestion(params, count),
+                    )
                 elif isinstance(incidents_result, dict):
                     count = incidents_result.get("total_count", len(incidents_result.get("data", [])))
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": incidents_result.get("data", []),
-                        "total_count": count,
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, count),
-                    }
+                    return ListRepoIncidentsResult(
+                        source_id=params.source_id,
+                        incidents=incidents_result.get("data", []),
+                        total_count=count,
+                        applied_filters=_build_filter_info(params),
+                        suggestion=_build_suggestion(params, count),
+                    )
                 else:
                     # Fallback for unexpected types
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": [],
-                        "total_count": 0,
-                        "error": f"Unexpected response type: {type(incidents_result).__name__}",
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, 0),
-                    }
+                    return ListRepoIncidentsError(
+                        error=f"Unexpected response type: {type(incidents_result).__name__}",
+                    )
             else:
                 incidents_result = await client.list_source_incidents(params.source_id, **api_params)
                 if isinstance(incidents_result, dict):
                     count = incidents_result.get("total_count", 0)
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": incidents_result.get("data", []),
-                        "next_cursor": incidents_result.get("next_cursor"),
-                        "total_count": count,
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, count),
-                    }
+                    return ListRepoIncidentsResult(
+                        source_id=params.source_id,
+                        incidents=incidents_result.get("data", []),
+                        next_cursor=incidents_result.get("next_cursor"),
+                        total_count=count,
+                        applied_filters=_build_filter_info(params),
+                        suggestion=_build_suggestion(params, count),
+                    )
                 elif isinstance(incidents_result, list):
                     # Handle case where API returns a list directly
                     count = len(incidents_result)
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": incidents_result,
-                        "total_count": count,
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, count),
-                    }
+                    return ListRepoIncidentsResult(
+                        source_id=params.source_id,
+                        incidents=incidents_result,
+                        total_count=count,
+                        applied_filters=_build_filter_info(params),
+                        suggestion=_build_suggestion(params, count),
+                    )
                 else:
                     # Fallback for unexpected types
-                    return {
-                        "source_id": params.source_id,
-                        "incidents": [],
-                        "total_count": 0,
-                        "error": f"Unexpected response type: {type(incidents_result).__name__}",
-                        "applied_filters": _build_filter_info(params),
-                        "suggestion": _build_suggestion(params, 0),
-                    }
+                    return ListRepoIncidentsError(
+                        error=f"Unexpected response type: {type(incidents_result).__name__}",
+                    )
         else:
             # Use repository_name lookup (legacy path)
             result = await client.list_repo_incidents_directly(
@@ -259,14 +264,20 @@ async def list_repo_incidents(params: ListRepoIncidentsParams) -> dict[str, Any]
                 mine=params.mine,
             )
 
-            # Enrich result with filter info
+            # Enrich result with filter info and convert to Pydantic model
             if isinstance(result, dict):
                 count = result.get("total_count", len(result.get("incidents", [])))
-                result["applied_filters"] = _build_filter_info(params)
-                result["suggestion"] = _build_suggestion(params, count)
-
-            return result
+                return ListRepoIncidentsResult(
+                    source_id=result.get("source_id"),
+                    incidents=result.get("incidents", []),
+                    total_count=count,
+                    next_cursor=result.get("next_cursor"),
+                    applied_filters=_build_filter_info(params),
+                    suggestion=_build_suggestion(params, count),
+                )
+            else:
+                return ListRepoIncidentsError(error="Unexpected result format from legacy path")
 
     except Exception as e:
         logger.error(f"Error listing repository incidents: {str(e)}")
-        return {"error": f"Failed to list repository incidents: {str(e)}"}
+        return ListRepoIncidentsError(error=f"Failed to list repository incidents: {str(e)}")

packages/gg_api_core/src/gg_api_core/tools/list_repo_occurrences.py

@@ -74,6 +74,22 @@ class ListRepoOccurrencesParams(ListRepoOccurrencesFilters, ListRepoOccurrencesB
     pass
 
 
+class ListRepoOccurrencesResult(BaseModel):
+    """Result from listing repository occurrences."""
+    repository: str | None = Field(default=None, description="Repository name")
+    occurrences_count: int = Field(description="Number of occurrences returned")
+    occurrences: list[dict[str, Any]] = Field(default_factory=list, description="List of occurrence objects")
+    cursor: str | None = Field(default=None, description="Pagination cursor for next page")
+    has_more: bool = Field(default=False, description="Whether more results are available")
+    applied_filters: dict[str, Any] = Field(default_factory=dict, description="Filters that were applied to the query")
+    suggestion: str = Field(default="", description="Suggestions for interpreting or modifying the results")
+
+
+class ListRepoOccurrencesError(BaseModel):
+    """Error result from listing repository occurrences."""
+    error: str = Field(description="Error message")
+
+
 def _build_filter_info(params: ListRepoOccurrencesParams) -> dict[str, Any]:
     """Build a dictionary describing the filters applied to the query."""
     filters = {}
@@ -127,7 +143,7 @@ def _build_suggestion(params: ListRepoOccurrencesParams, occurrences_count: int)
     return "\n".join(suggestions) if suggestions else ""
 
 
-async def list_repo_occurrences(params: ListRepoOccurrencesParams) -> dict[str, Any]:
+async def list_repo_occurrences(params: ListRepoOccurrencesParams) -> ListRepoOccurrencesResult | ListRepoOccurrencesError:
     """
     List secret occurrences for a specific repository using the GitGuardian v1/occurrences/secrets API.
 
@@ -159,7 +175,7 @@ async def list_repo_occurrences(params: ListRepoOccurrencesParams) -> dict[str,
 
     # Validate that at least one of repository_name or source_id is provided
     if not params.repository_name and not params.source_id:
-        return {"error": "Either repository_name or source_id must be provided"}
+        return ListRepoOccurrencesError(error="Either repository_name or source_id must be provided")
 
     logger.debug(f"Listing occurrences with repository_name={params.repository_name}, source_id={params.source_id}")
 
@@ -208,34 +224,34 @@ async def list_repo_occurrences(params: ListRepoOccurrencesParams) -> dict[str,
         if isinstance(result, dict):
             occurrences = result.get("occurrences", [])
             count = len(occurrences)
-            return {
-                "repository": params.repository_name,
-                "occurrences_count": count,
-                "occurrences": occurrences,
-                "cursor": result.get("cursor"),
-                "has_more": result.get("has_more", False),
-                "applied_filters": _build_filter_info(params),
-                "suggestion": _build_suggestion(params, count),
-            }
+            return ListRepoOccurrencesResult(
+                repository=params.repository_name,
+                occurrences_count=count,
+                occurrences=occurrences,
+                cursor=result.get("cursor"),
+                has_more=result.get("has_more", False),
+                applied_filters=_build_filter_info(params),
+                suggestion=_build_suggestion(params, count),
+            )
         elif isinstance(result, list):
             # If get_all=True, we get a list directly
             count = len(result)
-            return {
-                "repository": params.repository_name,
-                "occurrences_count": count,
-                "occurrences": result,
-                "applied_filters": _build_filter_info(params),
-                "suggestion": _build_suggestion(params, count),
-            }
+            return ListRepoOccurrencesResult(
+                repository=params.repository_name,
+                occurrences_count=count,
+                occurrences=result,
+                applied_filters=_build_filter_info(params),
+                suggestion=_build_suggestion(params, count),
+            )
         else:
-            return {
-                "repository": params.repository_name,
-                "occurrences_count": 0,
-                "occurrences": [],
-                "applied_filters": _build_filter_info(params),
-                "suggestion": _build_suggestion(params, 0),
-            }
+            return ListRepoOccurrencesResult(
+                repository=params.repository_name,
+                occurrences_count=0,
+                occurrences=[],
+                applied_filters=_build_filter_info(params),
+                suggestion=_build_suggestion(params, 0),
+            )
 
     except Exception as e:
         logger.error(f"Error listing repository occurrences: {str(e)}")
-        return {"error": f"Failed to list repository occurrences: {str(e)}"}
+        return ListRepoOccurrencesError(error=f"Failed to list repository occurrences: {str(e)}")

packages/gg_api_core/src/gg_api_core/tools/remediate_secret_incidents.py

@@ -56,7 +56,25 @@ def validate_source_or_repository(self) -> "RemediateSecretIncidentsParams":
         return self
 
 
-async def remediate_secret_incidents(params: RemediateSecretIncidentsParams) -> dict[str, Any]:
+class RemediateSecretIncidentsResult(BaseModel):
+    """Result from remediating secret incidents."""
+    repository_info: dict[str, Any] = Field(description="Information about the repository")
+    summary: dict[str, Any] | None = Field(default=None, description="Summary of occurrences, files, and secret types")
+    remediation_steps: list[dict[str, Any]] = Field(default_factory=list, description="Steps for remediating each file")
+    message: str | None = Field(default=None, description="Message when no occurrences found")
+    env_example_content: str | None = Field(default=None, description="Suggested .env.example content")
+    env_example_instructions: list[str] | None = Field(default=None, description="Instructions for .env.example")
+    git_commands: dict[str, Any] | None = Field(default=None, description="Git commands to fix history")
+    applied_filters: dict[str, Any] = Field(default_factory=dict, description="Filters applied when querying occurrences")
+    suggestion: str = Field(default="", description="Suggestions for interpreting results")
+
+
+class RemediateSecretIncidentsError(BaseModel):
+    """Error result from remediating secret incidents."""
+    error: str = Field(description="Error message")
+
+
+async def remediate_secret_incidents(params: RemediateSecretIncidentsParams) -> RemediateSecretIncidentsResult | RemediateSecretIncidentsError:
     """
     Find and remediate secret incidents in the current repository using EXACT match locations.
 
@@ -108,10 +126,17 @@ async def remediate_secret_incidents(params: RemediateSecretIncidentsParams) ->
         )
         occurrences_result = await list_repo_occurrences(occurrences_params)
 
-        if "error" in occurrences_result:
-            return {"error": occurrences_result["error"]}
+        # Check if list_repo_occurrences returned an error
+        if isinstance(occurrences_result, dict) and "error" in occurrences_result:
+            return RemediateSecretIncidentsError(error=occurrences_result["error"])
+
+        # Since list_repo_occurrences now returns Pydantic models, handle both old dict and new model formats
+        if hasattr(occurrences_result, 'model_dump'):
+            occurrences_dict = occurrences_result.model_dump()
+        else:
+            occurrences_dict = occurrences_result
 
-        occurrences = occurrences_result.get("occurrences", [])
+        occurrences = occurrences_dict.get("occurrences", [])
 
         # Filter by assignee if mine=True
         if params.mine:
@@ -132,13 +157,13 @@ async def remediate_secret_incidents(params: RemediateSecretIncidentsParams) ->
                 logger.warning(f"Could not filter by assignee: {str(e)}")
 
         if not occurrences:
-            return {
-                "repository_info": {"name": params.repository_name},
-                "message": "No secret occurrences found for this repository that match the criteria.",
-                "remediation_steps": [],
-                "applied_filters": occurrences_result.get("applied_filters", {}),
-                "suggestion": occurrences_result.get("suggestion", ""),
-            }
+            return RemediateSecretIncidentsResult(
+                repository_info={"name": params.repository_name},
+                message="No secret occurrences found for this repository that match the criteria.",
+                remediation_steps=[],
+                applied_filters=occurrences_dict.get("applied_filters", {}),
+                suggestion=occurrences_dict.get("suggestion", ""),
+            )
 
         # Process occurrences for remediation with exact location data
         logger.debug(f"Processing {len(occurrences)} occurrences with exact locations for remediation")
@@ -151,11 +176,12 @@ async def remediate_secret_incidents(params: RemediateSecretIncidentsParams) ->
         logger.debug(
             f"Remediation processing complete, returning result with {len(result.get('remediation_steps', []))} steps"
         )
-        return result
+        # Convert dict result to Pydantic model
+        return RemediateSecretIncidentsResult(**result)
 
     except Exception as e:
         logger.error(f"Error remediating incidents: {str(e)}")
-        return {"error": f"Failed to remediate incidents: {str(e)}"}
+        return RemediateSecretIncidentsError(error=f"Failed to remediate incidents: {str(e)}")
 
 
 async def _process_occurrences_for_remediation(