Update app.py

app.py

@@ -323,58 +323,67 @@ headers = {"Authorization": f"Bearer {qwen}"}
 
 def clean_output(output):
     """
-    Cleans the output from Qwen to ensure only required sections are displayed.
+    Cleans the output from the Hugging Face model to ensure only the relevant details are included.
     """
+    # Remove known markers for prompts and redundant sections
+    markers_to_exclude = [
+        "Functions identified by Gemini",
+        "Tasks:",
+        "Return only the required information",
+        "User-specified functionality:"
+    ]
+    
     lines = output.splitlines()
     filtered_lines = []
-    in_valid_section = False
-
+    seen_sections = set()
+    
     for line in lines:
         line = line.strip()
-        if line.startswith("Project Summary:") or line.startswith("Functionality Summary:") or line.startswith("Functionality Flow:"):
-            in_valid_section = True
-            filtered_lines.append(line)
-        elif line.startswith("Function Documentation:"):
-            in_valid_section = True
-            filtered_lines.append(line)
-        elif in_valid_section and line:
+        # Skip lines that are part of excluded markers
+        if any(marker in line for marker in markers_to_exclude):
+            continue
+        
+        # Skip duplicate sections
+        if line.startswith("Project Summary:") or line.startswith("Functionality Summary:"):
+            if line in seen_sections:
+                continue
+            seen_sections.add(line)
+        
+        # Skip empty lines
+        if line:
             filtered_lines.append(line)
-        elif line.startswith("File:") or line.startswith("User-specified functionality:"):
-            in_valid_section = False
-
+    
     return "\n".join(filtered_lines)
 
 
 
 
-def validate_and_generate_documentation(api_url, headers, gemini_output, functionality_description):
+
+def validate_and_generate_documentation(api_url, headers, gemini_output, file_contents, functionality_description):
     """
-    Uses the Hugging Face Inference API to generate clean and relevant documentation using Qwen.
+    Uses the Hugging Face Inference API to generate documentation in chunks to avoid token limits
+    and ensures only clean output is returned.
     """
-    # Clean Gemini output
-    cleaned_gemini_output = extract_cleaned_gemini_output(gemini_output)
-
-    # Generate the refined prompt for Qwen
-    prompt = f"""
+    # Generate the refined prompt for the Qwen model
+    base_prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
-    {cleaned_gemini_output}
-    
-    Tasks:
+    {gemini_output}
+    Qwen, identify the functions provided above in the project, and with the User-specified functionality in mind, perform these tasks:
     1. Generate a project summary:
         '
         Project Summary:
-        <Include project description and library or module dependencies>
+        <Qwen, include project description and library or module dependencies>
         '
     2. Refine the user-defined functionality:
         '
         Functionality Summary:
-        <Provide an enhanced description of user-specified functionality>
+        <Qwen, provide an enhanced description of user-specified functionality>
         '
     3. Describe the functionality flow:
         '
         Functionality Flow:
-        <Explain the sequence of functions and data flow>
+        <Qwen, explain the sequence of functions and data flow>
         '
     4. Generate detailed documentation for each function:
         '
@@ -390,20 +399,45 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, functio
           - Assumptions: <Any assumptions the function makes>
           - Example Usage: <Example demonstrating usage>
         '
-    5. Return only the required information for the above tasks, and exclude everything else.
     """
+    
+    # Split file contents into chunks to avoid exceeding token limits
+    max_chunk_size = 12000  # Adjust for tokenization overhead
+    file_chunks = []
+    current_chunk = base_prompt
+
+    for file_path, content in file_contents.items():
+        chunk_content = f"File: {os.path.basename(file_path)}\n{content}\n\n"
+        if len(current_chunk) + len(chunk_content) > max_chunk_size:
+            file_chunks.append(current_chunk)
+            current_chunk = base_prompt + chunk_content
+        else:
+            current_chunk += chunk_content
+
+    if current_chunk not in file_chunks:
+        file_chunks.append(current_chunk)
+
+    # Process each chunk through the API
+    full_output = ""
+    for chunk in file_chunks:
+        payload = {"inputs": chunk, "parameters": {"max_new_tokens": 2048}}
+        response = requests.post(api_url, headers=headers, json=payload)
+
+        if response.status_code == 200:
+            api_response = response.json()
+            if isinstance(api_response, list):
+                output = api_response[0].get("generated_text", "")
+            elif isinstance(api_response, dict):
+                output = api_response.get("generated_text", "")
+            else:
+                raise ValueError("Unexpected response format from Hugging Face API.")
+            full_output += output
+        else:
+            raise ValueError(f"Error during API call: {response.status_code}, {response.text}")
 
-    # Prepare payload and call API
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 1024}}
-    response = requests.post(api_url, headers=headers, json=payload)
+    # Apply cleaning to remove unnecessary content
+    return clean_output(full_output)
 
-    # Handle API response
-    if response.status_code == 200:
-        api_response = response.json()
-        output = api_response.get("generated_text", "") if isinstance(api_response, dict) else api_response[0].get("generated_text", "")
-        return clean_output(output)
-    else:
-        raise ValueError(f"Error during API call: {response.status_code}, {response.text}")
 
 
 
@@ -431,21 +465,30 @@ def generate_documentation_page():
                     # Call Gemini to identify required functions
                     gemini_result = identify_required_functions(project_folder, functionality)
 
-                    # Generate documentation using Qwen
+                    # Read project files
+                    file_paths = read_project_files(project_folder)
+                    file_contents = read_files(file_paths)
+
+                    # Call the Hugging Face API to generate documentation
                     final_documentation = validate_and_generate_documentation(
-                        API_URL, headers, gemini_result, functionality
+                        API_URL, headers, gemini_result, file_contents, functionality
                     )
 
-                    # Display the final documentation
+                    # Display the cleaned final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", final_documentation, height=600)
-                    
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
                 st.error("Project folder not found. Ensure the GitHub repository was cloned successfully.")
         else:
             st.error("Please enter the functionality to analyze.")
+
+    # Button to navigate back to the project page
+    if st.button("Back to Project"):
+        st.session_state.page = "project_view"
+        st.rerun()
+