diff --git a/tools/testgen.py b/tools/testgen.py
index 75b6ff2..b9aa372 100644
--- a/tools/testgen.py
+++ b/tools/testgen.py
@@ -418,7 +418,7 @@ class TestGenerationTool(BaseTool):
 
         return full_prompt
 
-    def format_response(self, response: str, request: TestGenRequest, model_info: Optional[dict] = None) -> str:
+    def format_response(self, response: str, request: TestGenerationRequest, model_info: Optional[dict] = None) -> str:
         """
         Format the test generation response.