Update response_formatter.py

response_formatter.py

@@ -1,12 +1,67 @@
+import re
+from typing import Optional
+
 from schemas import CodeTaskType, CodeXResponse, ResponseMeta
 
 
-def clean_text(text: str) -> str:
+def clean_text(text: Optional[str]) -> str:
     if not text:
         return ""
+    text = str(text).strip()
+    text = re.sub(r"\r\n", "\n", text)
+    text = re.sub(r"\n{3,}", "\n\n", text)
     return text.strip()
 
 
+def extract_section(text: str, label: str) -> Optional[str]:
+    pattern = rf"{label}:\s*(.*?)(?=\n[A-Z][A-Za-z ]*:\s|\Z)"
+    match = re.search(pattern, text, flags=re.DOTALL)
+    if match:
+        value = clean_text(match.group(1))
+        return value if value else None
+    return None
+
+
+def extract_code_block(text: str) -> Optional[str]:
+    fenced = re.findall(r"```(?:\w+)?\n(.*?)```", text, flags=re.DOTALL)
+    if fenced:
+        return clean_text(fenced[0])
+
+    code_section = extract_section(text, "Code")
+    if code_section:
+        return code_section
+
+    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 build_main_answer(text: str, explanation: Optional[str], code_output: Optional[str]) -> str:
+    if explanation:
+        return explanation
+
+    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)
+
+    if cleaned:
+        return cleaned
+
+    if code_output:
+        return "Code generated successfully."
+
+    return "Request processed successfully."
+
+
 def build_response(
     task_type: CodeTaskType,
     model_output: str,
@@ -14,15 +69,24 @@ def build_response(
     used_fallback: bool,
     retrieval_used: bool = False,
     source_count: int = 0,
+    processing_time_ms: Optional[int] = None,
 ) -> 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)
+
+    warnings = []
+    if not cleaned_output:
+        warnings.append("Model returned an empty response.")
+
     return CodeXResponse(
         task_type=task_type,
-        answer=cleaned_output,
-        code_output=None,
-        explanation=None,
-        warnings=[],
+        answer=answer,
+        code_output=code_output,
+        explanation=explanation,
+        warnings=warnings,
         sources=[],
         needs_clarification=False,
         meta=ResponseMeta(
@@ -30,7 +94,7 @@ def build_response(
             fallback_used=used_fallback,
             retrieval_used=retrieval_used,
             source_count=source_count,
-            processing_time_ms=None,
+            processing_time_ms=processing_time_ms,
         ),
     )
 
@@ -38,13 +102,16 @@ def build_response(
 def build_error_response(
     task_type: CodeTaskType,
     error_message: str,
+    processing_time_ms: Optional[int] = None,
 ) -> CodeXResponse:
+    cleaned_error = clean_text(error_message)
+
     return CodeXResponse(
         task_type=task_type,
-        answer="",
+        answer="Request processing failed.",
         code_output=None,
         explanation=None,
-        warnings=[clean_text(error_message)],
+        warnings=[cleaned_error] if cleaned_error else ["Unknown error occurred."],
         sources=[],
         needs_clarification=False,
         meta=ResponseMeta(
@@ -52,6 +119,6 @@ def build_error_response(
             fallback_used=False,
             retrieval_used=False,
             source_count=0,
-            processing_time_ms=None,
+            processing_time_ms=processing_time_ms,
         ),
     )