Improved model response handling to handle additional response statuses in future

Improved testgen; encourages follow-ups with less work in between and less token generation to avoid surpassing the 25K barrier
Improved coderevew tool to request a focused code review instead where a single-pass code review is too large or complex
This commit is contained in:
Fahad
2025-06-14 18:43:56 +04:00
parent ec5fee4409
commit 442decba70
8 changed files with 383 additions and 31 deletions

View File

@@ -37,7 +37,7 @@ from utils.conversation_memory import (
)
from utils.file_utils import read_file_content, read_files, translate_path_for_environment
from .models import ClarificationRequest, ContinuationOffer, ToolOutput
from .models import SPECIAL_STATUS_MODELS, ContinuationOffer, ToolOutput
logger = logging.getLogger(__name__)
@@ -1183,31 +1183,43 @@ When recommending searches, be specific about what information you need and why
logger = logging.getLogger(f"tools.{self.name}")
try:
# Try to parse as JSON to check for clarification requests
# Try to parse as JSON to check for special status requests
potential_json = json.loads(raw_text.strip())
if isinstance(potential_json, dict) and potential_json.get("status") == "clarification_required":
# Validate the clarification request structure
clarification = ClarificationRequest(**potential_json)
logger.debug(f"{self.name} tool requested clarification: {clarification.question}")
# Extract model information for metadata
metadata = {
"original_request": (request.model_dump() if hasattr(request, "model_dump") else str(request))
}
if model_info:
model_name = model_info.get("model_name")
if model_name:
metadata["model_used"] = model_name
if isinstance(potential_json, dict) and "status" in potential_json:
status_key = potential_json.get("status")
status_model = SPECIAL_STATUS_MODELS.get(status_key)
return ToolOutput(
status="clarification_required",
content=clarification.model_dump_json(),
content_type="json",
metadata=metadata,
)
if status_model:
try:
# Use Pydantic for robust validation of the special status
parsed_status = status_model.model_validate(potential_json)
logger.debug(f"{self.name} tool detected special status: {status_key}")
# Extract model information for metadata
metadata = {
"original_request": (
request.model_dump() if hasattr(request, "model_dump") else str(request)
)
}
if model_info:
model_name = model_info.get("model_name")
if model_name:
metadata["model_used"] = model_name
return ToolOutput(
status=status_key,
content=parsed_status.model_dump_json(),
content_type="json",
metadata=metadata,
)
except Exception as e:
# Invalid payload for known status, log warning and continue as normal response
logger.warning(f"Invalid {status_key} payload: {e}")
except (json.JSONDecodeError, ValueError, TypeError):
# Not a JSON clarification request, treat as normal response
# Not a JSON special status request, treat as normal response
pass
# Normal text response - format using tool-specific formatting

View File

@@ -36,6 +36,10 @@ class ToolOutput(BaseModel):
"success",
"error",
"clarification_required",
"full_codereview_required",
"focused_review_required",
"test_sample_needed",
"more_tests_required",
"resend_prompt",
"continuation_available",
] = "success"
@@ -50,6 +54,7 @@ class ToolOutput(BaseModel):
class ClarificationRequest(BaseModel):
"""Request for additional context or clarification"""
status: Literal["clarification_required"] = "clarification_required"
question: str = Field(..., description="Question to ask Claude for more context")
files_needed: Optional[list[str]] = Field(
default_factory=list, description="Specific files that are needed for analysis"
@@ -60,6 +65,48 @@ class ClarificationRequest(BaseModel):
)
class FullCodereviewRequired(BaseModel):
"""Request for full code review when scope is too large for quick review"""
status: Literal["full_codereview_required"] = "full_codereview_required"
important: Optional[str] = Field(None, description="Important message about escalation")
reason: Optional[str] = Field(None, description="Reason why full review is needed")
class FocusedReviewRequired(BaseModel):
"""Request for Claude to provide smaller, focused subsets of code for review"""
status: Literal["focused_review_required"] = "focused_review_required"
reason: str = Field(..., description="Why the current scope is too large for effective review")
suggestion: str = Field(
..., description="Suggested approach for breaking down the review into smaller, focused parts"
)
class TestSampleNeeded(BaseModel):
"""Request for additional test samples to determine testing framework"""
status: Literal["test_sample_needed"] = "test_sample_needed"
reason: str = Field(..., description="Reason why additional test samples are required")
class MoreTestsRequired(BaseModel):
"""Request for continuation to generate additional tests"""
status: Literal["more_tests_required"] = "more_tests_required"
pending_tests: str = Field(..., description="List of pending tests to be generated")
# Registry mapping status strings to their corresponding Pydantic models
SPECIAL_STATUS_MODELS = {
"clarification_required": ClarificationRequest,
"full_codereview_required": FullCodereviewRequired,
"focused_review_required": FocusedReviewRequired,
"test_sample_needed": TestSampleNeeded,
"more_tests_required": MoreTestsRequired,
}
class DiagnosticHypothesis(BaseModel):
"""A debugging hypothesis with context and next steps"""

View File

@@ -117,7 +117,10 @@ class TestGenTool(BaseTool):
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Can be used to continue conversations across different tools. Only provide this if continuing a previous conversation thread.",
"description": (
"Thread continuation ID for multi-turn conversations. Can be used to continue conversations "
"across different tools. Only provide this if continuing a previous conversation thread."
),
},
},
"required": ["files", "prompt"] + (["model"] if self.is_effective_auto_mode() else []),
@@ -436,7 +439,7 @@ class TestGenTool(BaseTool):
---
# IMMEDIATE ACTION REQUIRED
# IMMEDIATE NEXT ACTION
Claude, you are now in EXECUTION MODE. Take immediate action:
@@ -472,5 +475,6 @@ After creating each test file, show the user:
## Step 5: MOVE TO NEXT ACTION
Once tests are confirmed working, immediately proceed to the next logical step for the project.
**CRITICAL**: Do NOT stop after generating - you MUST create, validate, run, and confirm the tests work.
Take full ownership of the testing implementation and move to your next work."""
**CRITICAL**: Do NOT stop after generating - you MUST create, validate, run, and confirm the tests work. Take full
ownership of the testing implementation and move to your next work. If you were supplied a more_work_required request
in the response above, you MUST honor it."""