Update app.py

app.py

@@ -285,6 +285,31 @@ def identify_required_functions(project_path, functionality_description):
     return response.text
 
 
+def extract_cleaned_gemini_output(gemini_output):
+    """
+    Extracts and formats the cleaned output from Gemini to send to Qwen.
+    Args:
+        gemini_output (str): The output returned by Gemini.
+    Returns:
+        str: Cleaned and formatted output for Qwen.
+    """
+    lines = gemini_output.splitlines()
+    cleaned_output = []
+    functions_section = False
+
+    for line in lines:
+        line = line.strip()
+        if line.startswith("Project Summary:") or line.startswith("Functionality:"):
+            cleaned_output.append(line)
+        elif line.startswith("Functions:"):
+            cleaned_output.append(line)
+            functions_section = True
+        elif functions_section and line:
+            cleaned_output.append(line)
+        elif line.startswith("File:") or "Qwen," in line:
+            break
+
+    return "\n".join(cleaned_output)
 
 
 def split_into_chunks(content, chunk_size=1000):
@@ -298,120 +323,60 @@ headers = {"Authorization": f"Bearer {qwen}"}
 
 def clean_output(output):
     """
-    Cleans the output from the Hugging Face model to ensure no unnecessary details are included.
+    Cleans the output from Qwen to ensure only required sections are displayed.
     """
-    # Remove any remnants of prompts or file content (heuristically clean up known patterns)
     lines = output.splitlines()
-    filtered_lines = [
-        line for line in lines if not (
-            line.startswith("File:") or
-            line.startswith("User-specified functionality:") or
-            line.startswith("Functions identified by Gemini:") or
-            line.strip() == ""
-        )
-    ]
-    return "\n".join(filtered_lines)
-
-def extract_project_summary_and_functionality(gemini_output):
-    """
-    Extracts the Project Summary and Enhanced Functionality from the Gemini output.
-
-    Args:
-        gemini_output (str): The output returned by Gemini.
+    filtered_lines = []
+    in_valid_section = False
 
-    Returns:
-        dict: A dictionary with keys 'project_summary' and 'enhanced_functionality'.
-    """
-    lines = gemini_output.splitlines()
-    project_summary = []
-    enhanced_functionality = []
-    
-    # Extract relevant sections
-    current_section = None
     for line in lines:
         line = line.strip()
-        if line.startswith("Project Summary:"):
-            current_section = "project_summary"
-        elif line.startswith("Functionality:"):
-            current_section = "enhanced_functionality"
-        elif line.startswith("Functions:"):
-            break
-        elif current_section == "project_summary" and line:
-            project_summary.append(line)
-        elif current_section == "enhanced_functionality" and line:
-            enhanced_functionality.append(line)
+        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:
+            filtered_lines.append(line)
+        elif line.startswith("File:") or line.startswith("User-specified functionality:"):
+            in_valid_section = False
 
-    return {
-        "project_summary": " ".join(project_summary).replace("Project Summary:", "").strip(),
-        "enhanced_functionality": " ".join(enhanced_functionality).replace("Functionality:", "").strip(),
-    }
+    return "\n".join(filtered_lines)
 
-def extract_functions(gemini_output):
-    """
-    Extracts the functions and their dependencies from the Gemini output.
 
-    Args:
-        gemini_output (str): The output returned by Gemini.
 
-    Returns:
-        dict: A dictionary where keys are file names and values are dictionaries of functions and their dependencies.
-    """
-    lines = gemini_output.splitlines()
-    functions_section = False
-    functions = {}
-    current_file = None
 
-    for line in lines:
-        line = line.strip()
-        if line.startswith("Functions:"):
-            functions_section = True
-            continue
-        if functions_section:
-            # Identify file name
-            if line.endswith(":") and not line.startswith("-"):
-                current_file = line[:-1].strip()  # Remove colon and strip
-                functions[current_file] = {}
-            # Identify function details
-            elif line.startswith("-") and current_file:
-                parts = line.split(":", maxsplit=1)
-                if len(parts) == 2:
-                    function_header = parts[0].replace("-", "").strip()
-                    dependencies = parts[1].replace("Function Dependencies:", "").strip()
-                    functions[current_file][function_header] = dependencies.split(",") if dependencies else []
-
-    return functions
-
-
-def validate_and_generate_documentation(api_url, headers, gemini_output, file_contents, functionality_description):
+def validate_and_generate_documentation(api_url, headers, gemini_output, functionality_description):
     """
-    Uses the Hugging Face Inference API to generate documentation in chunks to avoid token limits.
+    Uses the Hugging Face Inference API to generate clean and relevant documentation using Qwen.
     """
-    # Generate the refined prompt for the Qwen model
-    base_prompt = f"""
+    # Clean Gemini output
+    cleaned_gemini_output = extract_cleaned_gemini_output(gemini_output)
+
+    # Generate the refined prompt for Qwen
+    prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
-    {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 summary of the project in this format:
+    {cleaned_gemini_output}
+    
+    Tasks:
+    1. Generate a project summary:
         '
         Project Summary:
-        <Qwen, include project description and library or module dependencies here>\n
+        <Include project description and library or module dependencies>
         '
-
-    2. Refine the user-defined functionality with your answer in this format:
+    2. Refine the user-defined functionality:
         '
         Functionality Summary:
-        <Qwen, provide an enhanced description of user-specified functionality here>\n
+        <Provide an enhanced description of user-specified functionality>
         '
-
-    3. Describe the flow of the functionality with your answer here:
+    3. Describe the functionality flow:
         '
         Functionality Flow:
-        <Qwen, Explain the sequence of functions and data flow>\n
+        <Explain the sequence of functions and data flow>
         '
-
-    4. For all relevant functions, generate detailed documentation in this format:
+    4. Generate detailed documentation for each function:
         '
         Function Documentation:
         For each relevant function:
@@ -420,50 +385,26 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, file_co
           - 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 in the function>
+          - 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 how to use the function>\n
+          - Example Usage: <Example demonstrating usage>
         '
-    5. Qwen, return only what was asked of you in the 4 tasks defined above, and nothing else
+    5. Return only the required information for the above tasks, and exclude everything else.
     """
 
-    # Split file contents into chunks to avoid exceeding the token limit
-    max_chunk_size = 12000  # Adjust based on the tokenization overhead
-    file_chunks = []
-    current_chunk = base_prompt
+    # Prepare payload and call API
+    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 1024}}
+    response = requests.post(api_url, headers=headers, json=payload)
 
-    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
-
-    # Add the last chunk
-    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": 1024}}
-        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}")
+    # 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}")
 
-    return full_output
 
 
 def generate_documentation_page():
@@ -490,29 +431,22 @@ def generate_documentation_page():
                     # Call Gemini to identify required functions
                     gemini_result = identify_required_functions(project_folder, functionality)
 
-                    # Read project files
-                    file_paths = read_project_files(project_folder)
-                    file_contents = read_files(file_paths)
-
-                    # Call the Hugging Face API to generate documentation
+                    # Generate documentation using Qwen
                     final_documentation = validate_and_generate_documentation(
-                        API_URL, headers, gemini_result, file_contents, functionality
+                        API_URL, headers, gemini_result, functionality
                     )
 
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", final_documentation, height=600)
+                    st.rerun()
                 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()
+