Update app.py

app.py

@@ -391,8 +391,8 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, functio
     """
     Uses the Hugging Face Inference API to generate clean and relevant documentation using Qwen.
     """
+    # Clean Gemini output
     cleaned_gemini_output = extract_cleaned_gemini_output(gemini_output)
-    gemini_output_length = len(cleaned_gemini_output)  # Record the length of the cleaned Gemini output
 
     # Generate the refined prompt for Qwen
     prompt = f"""
@@ -433,8 +433,6 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, functio
     5. Return only the required information for the above tasks, and exclude everything else.
     """
 
-    qwen_prompt_length = len(prompt)
-    
     payload = {"inputs": prompt, "parameters": {"max_new_tokens": 1024}}
     response = requests.post(api_url, headers=headers, json=payload)
 
@@ -443,14 +441,23 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, functio
         api_response = response.json()
         output = api_response.get("generated_text", "") if isinstance(api_response, dict) else api_response[0].get("generated_text", "")
 
-        # Remove the Gemini content from the top of the Qwen output
-        trimmed_output = output[gemini_output_length + qwen_prompt_length:].strip()
+        # Extract only the relevant sections by identifying key markers in the output
+        relevant_sections = []
+        capturing = False
+
+        for line in output.splitlines():
+            if "Project Summary:" in line or "Functionality Summary:" in line or "Functionality Flow:" in line or "Function Documentation:" in line:
+                capturing = True  # Start capturing from relevant sections
+            if capturing:
+                relevant_sections.append(line)
 
-        return clean_output(trimmed_output)  # Final cleanup if necessary
+        # Join relevant sections and perform final cleanup
+        return clean_output("\n".join(relevant_sections))
     else:
         raise ValueError(f"Error during API call: {response.status_code}, {response.text}")
 
 
+
 def generate_documentation_page():
     st.subheader(f"Generate Documentation for {st.session_state.current_project}")
     st.write("Enter the functionality or parts of the project for which you'd like to generate documentation.")