update code

response_formatter.py

@@ -34,34 +34,126 @@ def extract_code_block(text: str) -> Optional[str]:
     return None
 
 
-def extract_explanation(text: str) -> Optional[str]:
-    for label in ["Explanation", "Root Cause", "Review", "Suggestions"]:
-        value = extract_section(text, label)
-        if value:
-            return value
-    return None
+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")
+
+    if task_type == CodeTaskType.REVIEW:
+        review = extract_section(text, "Review")
+        suggestions = extract_section(text, "Suggestions")
+        if review and suggestions:
+            return f"{review}\n\nSuggestions:\n{suggestions}"
+        return review or suggestions
+
+    if task_type == CodeTaskType.REFACTOR:
+        return extract_section(text, "Explanation")
+
+    if task_type == CodeTaskType.EXPLAIN:
+        return extract_section(text, "Explanation")
+
+    if task_type == CodeTaskType.GENERATE:
+        return extract_section(text, "Explanation")
+
+    return (
+        extract_section(text, "Explanation")
+        or extract_section(text, "Root Cause")
+        or extract_section(text, "Review")
+        or extract_section(text, "Suggestions")
+    )
 
 
-def build_main_answer(text: str, explanation: Optional[str], code_output: Optional[str]) -> str:
-    if explanation:
-        return explanation
+def strip_structured_sections(text: str) -> str:
+    text = re.sub(r"```(?:\w+)?\n.*?```", "", text, flags=re.DOTALL)
+    text = re.sub(r"Code:\s*.*", "", text, flags=re.DOTALL)
+    text = re.sub(r"Explanation:\s*.*", "", text, flags=re.DOTALL)
+    text = re.sub(r"Root Cause:\s*.*", "", text, flags=re.DOTALL)
+    text = re.sub(r"Review:\s*.*", "", text, flags=re.DOTALL)
+    text = re.sub(r"Suggestions:\s*.*", "", text, flags=re.DOTALL)
+    return clean_text(text)
 
-    cleaned = clean_text(text)
 
-    if code_output:
-        cleaned = re.sub(r"```(?:\w+)?\n.*?```", "", cleaned, flags=re.DOTALL)
-        cleaned = re.sub(r"Code:\s*.*", "", cleaned, flags=re.DOTALL)
-        cleaned = clean_text(cleaned)
+def first_sentence(text: str) -> str:
+    text = clean_text(text)
+    if not text:
+        return ""
+    match = re.split(r"(?<=[.!?])\s+", text, maxsplit=1)
+    return clean_text(match[0]) if match else text
+
 
-    if cleaned:
-        return cleaned
+def build_main_answer(
+    task_type: CodeTaskType,
+    raw_text: str,
+    explanation: Optional[str],
+    code_output: Optional[str],
+) -> str:
+    cleaned = clean_text(raw_text)
+
+    if task_type == CodeTaskType.GENERATE:
+        if explanation:
+            return first_sentence(explanation)
+        if code_output:
+            return "Code generated successfully."
+        return "Generation completed."
+
+    if task_type == CodeTaskType.EXPLAIN:
+        if explanation:
+            return first_sentence(explanation)
+        stripped = strip_structured_sections(cleaned)
+        if stripped:
+            return first_sentence(stripped)
+        return "Code explanation generated."
+
+    if task_type == CodeTaskType.FIX:
+        root_cause = extract_section(cleaned, "Root Cause")
+        if root_cause:
+            return first_sentence(root_cause)
+        if explanation:
+            return first_sentence(explanation)
+        if code_output:
+            return "Code fix generated successfully."
+        return "Fix completed."
+
+    if task_type == CodeTaskType.REVIEW:
+        review = extract_section(cleaned, "Review")
+        if review:
+            return first_sentence(review)
+        if explanation:
+            return first_sentence(explanation)
+        stripped = strip_structured_sections(cleaned)
+        if stripped:
+            return first_sentence(stripped)
+        return "Code review completed."
+
+    if task_type == CodeTaskType.REFACTOR:
+        if explanation:
+            return first_sentence(explanation)
+        if code_output:
+            return "Code refactored successfully."
+        return "Refactor completed."
+
+    stripped = strip_structured_sections(cleaned)
+    if stripped:
+        return first_sentence(stripped)
+
+    if explanation:
+        return first_sentence(explanation)
 
     if code_output:
-        return "Code generated successfully."
+        return "Request processed successfully."
 
     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,
@@ -73,9 +165,10 @@ def build_response(
 ) -> CodeXResponse:
     cleaned_output = clean_text(model_output)
 
-    code_output = extract_code_block(cleaned_output)
-    explanation = extract_explanation(cleaned_output)
-    answer = build_main_answer(cleaned_output, explanation, code_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)
+    answer = build_main_answer(task_type, cleaned_output, explanation, code_output)
 
     warnings = []
     if not cleaned_output: