Update app.py

app.py

@@ -390,11 +390,31 @@ def clean_output(output):
 def validate_and_generate_documentation(api_url, headers, gemini_output, functionality_description):
     """
     Uses the Hugging Face Inference API to generate clean and relevant documentation using Qwen.
+    Handles truncated output by identifying missing functions and re-prompting Qwen.
     """
-    # Clean Gemini output
+    def get_missing_functions(output, cleaned_gemini_output):
+        """
+        Identifies functions mentioned in Gemini's output that are missing in Qwen's response.
+        """
+        gemini_functions = set()
+        for line in cleaned_gemini_output.splitlines():
+            if line.startswith("-public "):  # Assuming function definitions start with "-public"
+                function_signature = line.split(":")[0]
+                gemini_functions.add(function_signature.strip())
+        
+        qwen_functions = set()
+        for line in output.splitlines():
+            if line.startswith("- **Function**:"):  # Qwen's format for functions
+                function_signature = line.replace("- **Function**:", "").strip().split("`")[1]
+                qwen_functions.add(function_signature.strip())
+        
+        # Find functions in Gemini's output but missing in Qwen's response
+        return gemini_functions - qwen_functions
+
     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
+    # Generate the initial refined prompt for Qwen
     prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
@@ -433,6 +453,8 @@ 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)
 
@@ -441,23 +463,49 @@ 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", "")
 
-        # 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)
-
-        # Join relevant sections and perform final cleanup
-        return clean_output("\n".join(relevant_sections))
+        # Remove the Gemini content from the top of the Qwen output
+        trimmed_output = output[gemini_output_length + qwen_prompt_length:].strip()
+
+        # Check for missing functions
+        missing_functions = get_missing_functions(trimmed_output, cleaned_gemini_output)
+
+        if missing_functions:
+            # Re-prompt Qwen for the missing functions
+            missing_functions_prompt = f"""
+            User-specified functionality: '{functionality_description}'
+            The following functions were not fully documented in the previous response:
+            {', '.join(missing_functions)}
+            
+            Please provide detailed documentation for these functions in the following format:
+            '
+            Function Documentation:
+            For each function:
+              - Summary: <Description of the function's purpose>
+              - Inputs: <Details of inputs and their types>
+              - Outputs: <Details of outputs and their types>
+              - Dependencies: <Dependencies on other modules/functions>
+              - Data structures: <Details of data structures used>
+              - Algorithmic Details: <Description of the algorithm used>
+              - Error Handling: <Description of how the function handles errors>
+              - Assumptions: <Any assumptions the function makes>
+              - Example Usage: <Example demonstrating usage>
+            '
+            """
+            missing_payload = {"inputs": missing_functions_prompt, "parameters": {"max_new_tokens": 1024}}
+            missing_response = requests.post(api_url, headers=headers, json=missing_payload)
+
+            if missing_response.status_code == 200:
+                missing_api_response = missing_response.json()
+                missing_output = missing_api_response.get("generated_text", "") if isinstance(missing_api_response, dict) else missing_api_response[0].get("generated_text", "")
+                trimmed_output += "\n\n" + missing_output
+
+        return clean_output(trimmed_output)  # Final cleanup if necessary
     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.")