ui update

app.py

@@ -7,28 +7,39 @@ agent = TranslationAgent()
 title = "Unity → Unreal Code Translator"
 description = "Paste a Unity C# script and get Unreal C++ (.h/.cpp) + Unreal Blueprint Python generator."
 
+
 def extract_class_name(csharp_code: str) -> str:
     import re
     match = re.search(r"class\s+(\w+)", csharp_code)
     return match.group(1) if match else "MyActor"
 
+
+def split_cpp_code(full_cpp_code: str):
+    """Split combined C++ output into .h and .cpp sections."""
+    if not full_cpp_code:
+        return "", ""
+    lines = full_cpp_code.splitlines()
+    header_lines, cpp_lines = [], []
+    in_header = True
+    for line in lines:
+        if line.strip().startswith("#include") or line.strip().endswith(".cpp"):
+            in_header = False
+        if in_header:
+            header_lines.append(line)
+        else:
+            cpp_lines.append(line)
+    return "\n".join(header_lines).strip(), "\n".join(cpp_lines).strip()
+
+
 def do_translate(csharp_code, target_format, model_choice):
     """Run translation and return all outputs."""
     class_name = extract_class_name(csharp_code)
-    outputs = agent.translate(csharp_code, target_format=target_format, model_choice=model_choice)
+    cpp_code, blueprint_code = agent.translate(
+        csharp_code, target_format=target_format, model_choice=model_choice
+    )
 
-    # Handle tuple or dict return from the agent
-    if isinstance(outputs, tuple):
-        if len(outputs) == 3:
-            header_code, cpp_code, blueprint_code = outputs
-        else:
-            header_code, cpp_code, blueprint_code = outputs[0], "", ""
-    elif isinstance(outputs, dict):
-        header_code = outputs.get("header_code", "")
-        cpp_code = outputs.get("cpp_code", "")
-        blueprint_code = outputs.get("blueprint_code", "")
-    else:
-        header_code = cpp_code = blueprint_code = ""
+    # Split .h / .cpp parts if present
+    header_code, cpp_code_only = split_cpp_code(cpp_code)
 
     # Save files
     h_filename = f"{class_name}.h"
@@ -36,14 +47,14 @@ def do_translate(csharp_code, target_format, model_choice):
     with open(h_filename, "w", encoding="utf-8") as h:
         h.write(header_code or "// No header generated")
     with open(cpp_filename, "w", encoding="utf-8") as c:
-        c.write(cpp_code or "// No cpp generated")
+        c.write(cpp_code_only or "// No cpp generated")
 
-    # Split blueprint output if it has explanations/logs
+    # Split blueprint output into main and logs if available
     blueprint_parts = blueprint_code.split("###")
     main_script = blueprint_parts[0].strip()
     extra_info = blueprint_parts[1].strip() if len(blueprint_parts) > 1 else ""
 
-    return header_code, cpp_code, main_script, extra_info, h_filename, cpp_filename
+    return header_code, cpp_code_only, main_script, extra_info, h_filename, cpp_filename
 
 
 with gr.Blocks(title=title) as demo:
@@ -76,7 +87,8 @@ with gr.Blocks(title=title) as demo:
         with gr.Row():
             blueprint_main = gr.Textbox(label="Blueprint Python Script", lines=15)
         with gr.Row():
-            blueprint_notes = gr.Textbox(label="Blueprint Hints / Logs", lines=10, placeholder="Any additional output, logs, or notes will appear here.")
+            blueprint_notes = gr.Textbox(label="Blueprint Hints / Logs", lines=10,
+                                         placeholder="Any additional output, logs, or notes will appear here.")
 
     translate_btn.click(
         do_translate,