Update code

response_formatter.py

@@ -36,7 +36,11 @@ def extract_code_block(text: str) -> Optional[str]:
 
 def extract_explanation(text: str, task_type: CodeTaskType) -> Optional[str]:
     if task_type == CodeTaskType.FIX:
-        return extract_section(text, "Explanation") or extract_section(text, "Root Cause")
+        root_cause = extract_section(text, "Root Cause")
+        explanation = extract_section(text, "Explanation")
+        if root_cause and explanation:
+            return explanation
+        return explanation or root_cause
 
     if task_type == CodeTaskType.REVIEW:
         review = extract_section(text, "Review")
@@ -122,6 +126,78 @@ def build_review_summary(review_text: str) -> str:
     return "Code review completed."
 
 
+def extract_original_callable_name(raw_text: str) -> Optional[str]:
+    patterns = [
+        r"\bdef\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(",
+        r"\bclass\s+([A-Za-z_][A-Za-z0-9_]*)\s*[:(]",
+        r"\bfunction\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(",
+    ]
+    for pattern in patterns:
+        match = re.search(pattern, raw_text)
+        if match:
+            return match.group(1)
+    return None
+
+
+def replace_callable_name(code_output: str, old_name: str, new_name: str) -> str:
+    if not code_output or not old_name or not new_name or old_name == new_name:
+        return code_output
+
+    patterns = [
+        (rf"(\bdef\s+){re.escape(old_name)}(\s*\()", rf"\1{new_name}\2"),
+        (rf"(\bclass\s+){re.escape(old_name)}(\b)", rf"\1{new_name}\2"),
+        (rf"(\bfunction\s+){re.escape(old_name)}(\s*\()", rf"\1{new_name}\2"),
+    ]
+
+    updated = code_output
+    for pattern, replacement in patterns:
+        updated = re.sub(pattern, replacement, updated, count=1)
+    return updated
+
+
+def align_refactor_names_with_input(raw_text: str, code_output: Optional[str]) -> Optional[str]:
+    if not code_output:
+        return None
+
+    original_name = extract_original_callable_name(raw_text)
+    new_name = extract_original_callable_name(code_output)
+
+    if original_name and new_name and original_name != new_name:
+        return replace_callable_name(code_output, new_name, original_name)
+
+    return code_output
+
+
+def should_include_code_output(
+    task_type: CodeTaskType,
+    code_output: Optional[str],
+    explanation: Optional[str],
+) -> Optional[str]:
+    if not code_output:
+        return None
+
+    if task_type == CodeTaskType.EXPLAIN:
+        return None
+
+    if task_type == CodeTaskType.REVIEW:
+        if explanation:
+            lowered = explanation.lower()
+            review_only_signals = [
+                "consider",
+                "suggest",
+                "add a check",
+                "type hints",
+                "docstring",
+                "could improve",
+                "readability",
+                "maintainability",
+            ]
+            if any(signal in lowered for signal in review_only_signals):
+                return None
+
+    return code_output
+
+
 def build_main_answer(
     task_type: CodeTaskType,
     raw_text: str,
@@ -186,16 +262,6 @@ def build_main_answer(
     return "Request processed successfully."
 
 
-def should_include_code_output(task_type: CodeTaskType, code_output: Optional[str]) -> Optional[str]:
-    if not code_output:
-        return None
-
-    if task_type == CodeTaskType.EXPLAIN:
-        return None
-
-    return code_output
-
-
 def build_response(
     task_type: CodeTaskType,
     model_output: str,
@@ -204,18 +270,29 @@ def build_response(
     retrieval_used: bool = False,
     source_count: int = 0,
     processing_time_ms: Optional[int] = None,
+    original_code: Optional[str] = None,
 ) -> CodeXResponse:
     cleaned_output = clean_text(model_output)
 
     raw_code_output = extract_code_block(cleaned_output)
     explanation = extract_explanation(cleaned_output, task_type)
-    code_output = should_include_code_output(task_type, raw_code_output)
+
+    if task_type == CodeTaskType.REFACTOR and original_code:
+        raw_code_output = align_refactor_names_with_input(original_code, raw_code_output)
+
+    code_output = should_include_code_output(task_type, raw_code_output, explanation)
     answer = build_main_answer(task_type, cleaned_output, explanation, code_output)
 
     warnings = []
     if not cleaned_output:
         warnings.append("Model returned an empty response.")
 
+    if task_type == CodeTaskType.REFACTOR and original_code and raw_code_output and code_output:
+        original_name = extract_original_callable_name(original_code)
+        final_name = extract_original_callable_name(code_output)
+        if original_name and final_name and original_name != final_name:
+            warnings.append("Refactor output may have changed original naming unexpectedly.")
+
     return CodeXResponse(
         task_type=task_type,
         answer=answer,