Update app.py

app.py

@@ -246,12 +246,24 @@ def generate_prompt(file_contents, functionality_description):
     for file_path, content in file_contents.items():
         prompt += f"File: {os.path.basename(file_path)}\n{content}\n\n"
     
-    prompt += "If the functionality does not exist in the project, return 'Functionality does not exist.', otherwise, for each relevant function, provide:\n"
+    prompt += "For each relevant function, provide:\n"
     prompt += "1. Which file the function is found in.\n"
     prompt += "2. The function name.\n"
     prompt += "3. Dependencies on other functions or modules.\n"
 
-    prompt += "Also, return a more clear version of the user entered functionality at the top of your response. Provide no commentary, just what has been asked please."
+    prompt += """ 
+            Return your output in the following format providing no commentary:
+
+            Project Summary:
+                <A summary of the project>
+            Functionality:
+                <a more precise version of the user defined funcionality>
+            Functions:
+                <for each file containing relevant functions>
+                    <file name of the file that contains relevant functions>:
+                        -<relevant function header from the current file>:
+                            -Function Dependencies:<function dependencies seperated by commas>
+    """
 
     return prompt
 
@@ -300,42 +312,112 @@ def clean_output(output):
     ]
     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.
+
+    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)
+
+    return {
+        "project_summary": " ".join(project_summary).replace("Project Summary:", "").strip(),
+        "enhanced_functionality": " ".join(enhanced_functionality).replace("Functionality:", "").strip(),
+    }
+
+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):
     """
-    Uses the Hugging Face Inference API to validate functions and generate documentation.
-    Ensures that the output excludes file content and internal prompts.
+    Uses the Hugging Face Inference API to generate documentation based on Gemini output.
+    Ensures that the output matches the expected format without unnecessary details.
     """
-    # Generate the prompt for the Qwen model
+    # Generate the refined prompt for the Qwen model
     prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
     {gemini_output}
 
     Task:
-    1. Validate if the functions identified by Gemini are sufficient for implementing the functionality.
-    2. If not, identify all additional functions required.
-    3. For all relevant functions, generate detailed documentation in this format:
-    
-    Project Summary:
-    <Include project description and dependencies>
-
-    Functionality Summary:
-    <Description of user-specified functionality>
-
-    Functionality Flow:
-    <Explain the sequence of functions and data flow>
-
-    Function Documentation:
-    For each relevant 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 in the function>
-      - Error Handling: <Description of how the function handles errors>
-      - Assumptions: <Any assumptions the function makes>
-      - Example Usage: <Example demonstrating how to use the function>
+    1. Generate a summary of the project:
+        Project Summary:
+        <Include project description and library or module dependencies>
+
+    2. Refine the user-defined functionality:
+        Functionality Summary:
+        <Description of user-specified functionality>
+
+    3. Describe the flow of the functionality:
+        Functionality Flow:
+        <Explain the sequence of functions and data flow>
+
+    4. For all relevant functions, generate detailed documentation:
+        Function Documentation:
+        For each relevant 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 in the function>
+          - Error Handling: <Description of how the function handles errors>
+          - Assumptions: <Any assumptions the function makes>
+          - Example Usage: <Example demonstrating how to use the function>
     """
 
     # Send the prompt to the Hugging Face API
@@ -352,45 +434,48 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, file_co
         else:
             raise ValueError("Unexpected response format from Hugging Face API.")
         
-        # Clean the output to exclude any file contents or prompts
-        cleaned_output = clean_output(output)
-        return cleaned_output
+        # Clean the output to exclude any unnecessary details
+        return output
     else:
         raise ValueError(f"Error: {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 identify relevant functions.")
+    st.write("Enter the functionality or parts of the project for which you'd like to generate documentation.")
 
     # Prompt user for functionality description
-    functionality = st.text_area("Describe the functionality", placeholder="e.g., Explain the function of the file `main.py`")
+    functionality = st.text_area(
+        "Describe the functionality",
+        placeholder="e.g., Explain the function of the file `main.py`",
+    )
 
     # Button to start analyzing functionality
     if st.button("Analyze"):
         if functionality.strip():
             st.write("Analyzing project files... Please wait.")
-            
+
             # Get the path of the current project
             user_folder = os.path.join("user_projects", st.session_state.username)
             project_folder = os.path.join(user_folder, st.session_state.current_project)
-            
+
             if os.path.exists(project_folder):
                 try:
                     # 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 for validation and documentation
+
+                    # Call the Hugging Face API to generate documentation
                     final_documentation = validate_and_generate_documentation(
                         API_URL, headers, gemini_result, file_contents, functionality
                     )
-                    
+
                     # Display the final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", final_documentation, height=600)