Update app.py

app.py

@@ -324,102 +324,20 @@ def split_into_chunks(content, chunk_size=1000):
     """Splits large content into smaller chunks."""
     return [content[i:i + chunk_size] for i in range(0, len(content), chunk_size)]
 
-# Hugging Face Inference API endpoint for the model
-API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B-Instruct"
-qwen = os.getenv("QWEN")
-headers = {"Authorization": f"Bearer {qwen}"}
-
-def extract_cleaned_gemini_output(gemini_output):
+def generate_detailed_documentation(file_contents, functionality_description):
     """
-    Extracts and formats the cleaned output from Gemini to send to Qwen.
+    Generates detailed documentation using Gemini directly.
     Args:
-        gemini_output (str): The output returned by Gemini.
+        file_contents (dict): A dictionary with file paths as keys and their content as values.
+        functionality_description (str): A description of the functionality to document.
     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 clean_output(output):
-    """
-    Cleans the final output to include only documentation sections.
+        str: The generated documentation.
     """
-    lines = output.splitlines()
-    filtered_lines = []
-    in_valid_section = False
-
-    for line in lines:
-        line = line.strip()
-        # Include only valid sections
-        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)
-        # Stop processing at any undesired section
-        elif line.startswith("User-specified functionality:") or line.startswith("Tasks identified by Gemini:"):
-            in_valid_section = False
-        elif in_valid_section and line:
-            filtered_lines.append(line)
-
-    return "\n".join(filtered_lines)
-
-
-
-
+    prompt = f"""
+    The following code files are provided. Analyze their contents and generate comprehensive documentation.
 
+    Functionality description: '{functionality_description}'
 
-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.
-    """
-    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 initial refined prompt for Qwen
-    prompt = f"""
-    User-specified functionality: '{functionality_description}'
-    Functions identified by Gemini:
-    {cleaned_gemini_output}
-    
     Tasks:
     1. Generate a project summary:
         '
@@ -436,10 +354,10 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, functio
         Functionality Flow:
         <Explain the sequence of functions and data flow>
         '
-    4. Generate detailed documentation for each function:
+    4. Generate detailed documentation for each function in the codebase:
         '
         Function Documentation:
-        For each relevant function:
+        For each function:
           - Summary: <Description of the function's purpose>
           - Inputs: <Details of inputs and their types>
           - Outputs: <Details of outputs and their types>
@@ -450,60 +368,15 @@ 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.
-    """
-
-    qwen_prompt_length = len(prompt)
-    
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 1024}}
-    response = requests.post(api_url, headers=headers, json=payload)
-
-    # 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", "")
-
-        # 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}")
+    Please return only the required documentation in the specified format.
 
+    Code files:
+    """
+    for file_path, content in file_contents.items():
+        prompt += f"\nFile: {os.path.basename(file_path)}\n{content}\n"
 
+    response = model.generate_content(prompt)
+    return response.text.strip()
 
 
 def generate_documentation_page():
@@ -527,17 +400,18 @@ def generate_documentation_page():
 
             if os.path.exists(project_folder):
                 try:
-                    # Call Gemini to identify required functions
-                    gemini_result = identify_required_functions(project_folder, functionality)
+                    # Gather all file paths in the project directory
+                    file_paths = read_project_files(project_folder)
+
+                    # Read file contents
+                    file_contents = read_files(file_paths)
 
-                    # Generate documentation using Qwen
-                    final_documentation = validate_and_generate_documentation(
-                        API_URL, headers, gemini_result, functionality
-                    )
+                    # Generate documentation using Gemini
+                    documentation = generate_detailed_documentation(file_contents, functionality)
 
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
-                    st.text_area("Generated Documentation", final_documentation, height=600)
+                    st.text_area("Generated Documentation", documentation, height=600)
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
@@ -545,14 +419,6 @@ def generate_documentation_page():
         else:
             st.error("Please enter the functionality to analyze.")
 
-
-
-        
-
-
-
-
-
 #------------------------------------------------------------------------------------------------------------------------------------------------------------------------