Update app.py

app.py

@@ -321,52 +321,69 @@ API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B-In
 qwen = os.getenv("QWEN")
 headers = {"Authorization": f"Bearer {qwen}"}
 
-def clean_output(output):
+def extract_cleaned_gemini_output(gemini_output):
     """
-    Cleans the output from the model to ensure only relevant content is included.
-    Strips away any redundant prompts, instructions, and markers.
+    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 = 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.startswith("Qwen,") or
-            line.startswith("<") or
-            line.startswith("Tasks:") or
-            line.startswith("'") or
-            line.strip() == ""  # Exclude empty or irrelevant lines
-        )
-    ]
-    return "\n".join(filtered_lines)
+    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 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):
     """
-    Generates documentation by communicating with the Qwen model in manageable chunks.
-    Cleans the output to ensure user sees only relevant information.
+    Uses the Hugging Face Inference API to generate clean and relevant documentation using Qwen.
     """
-    # Restore the detailed prompt for Qwen
-    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
-    2. Refine the user-defined functionality with your answer in this format:
+        <Include project description and library or module dependencies>
+        '
+    2. Refine the user-defined functionality:
+        '
         Functionality Summary:
-        <Qwen, provide an enhanced description of user-specified functionality here>\n
-    3. Describe the flow of the functionality with your answer here:
+        <Provide an enhanced description of user-specified functionality>
+        '
+    3. Describe the functionality flow:
+        '
         Functionality Flow:
-        <Qwen, Explain the sequence of functions and data flow>\n
-    4. For all relevant functions, generate detailed documentation in this format:
+        <Explain the sequence of functions and data flow>
+        '
+    4. Generate detailed documentation for each function:
+        '
         Function Documentation:
         For each relevant function:
           - Summary: <Description of the function's purpose>
@@ -374,50 +391,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
-    Qwen, return only what was asked of you in the 4 tasks defined above, and nothing else
+          - Example Usage: <Example demonstrating usage>
+        '
+    5. Return only the required information for the above tasks, and exclude everything else.
     """
 
-    # Split file contents into manageable chunks
-    max_chunk_size = 12000  # Adjust for 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 final chunk
-    if current_chunk not in file_chunks:
-        file_chunks.append(current_chunk)
-
-    # Process each chunk and accumulate the cleaned output
-    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 += clean_output(output)  # Clean each chunk's 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():
     st.subheader(f"Generate Documentation for {st.session_state.current_project}")
@@ -443,16 +436,12 @@ 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 cleaned documentation
+                    # Display the final documentation
                     st.success("Documentation generated successfully!")
                     st.text_area("Generated Documentation", final_documentation, height=600)
                 except Exception as e:
@@ -462,10 +451,6 @@ def generate_documentation_page():
         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()