Update app.py

app.py

@@ -284,23 +284,33 @@ 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):
+    """
+    Cleans the output from the Hugging Face model to ensure no unnecessary details are included.
+    """
+    # 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 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."""
-    # Generate the base prompt for the Qwen model
+    """
+    Uses the Hugging Face Inference API to validate functions and generate documentation.
+    Ensures that the output excludes file content and internal prompts.
+    """
+    # Generate the prompt for the Qwen model
     prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
     {gemini_output}
 
-    Project files:
-    """
-    for file_path, content in file_contents.items():
-        # Split content into manageable chunks if necessary
-        chunks = split_into_chunks(content, chunk_size=3000)  # Adjust chunk size based on model's max input length
-        for idx, chunk in enumerate(chunks):
-            prompt += f"File: {os.path.basename(file_path)} (Part {idx + 1})\n{chunk}\n\n"
-
-    prompt += """
     Task:
     1. Validate if the functions identified by Gemini are sufficient for implementing the functionality.
     2. If not, identify all additional functions required.
@@ -322,6 +332,10 @@ 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>
+      - 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
@@ -332,15 +346,21 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, file_co
     if response.status_code == 200:
         api_response = response.json()
         if isinstance(api_response, list):  # If response is a list, extract the first element
-            return api_response[0].get("generated_text", "No output generated.")
+            output = api_response[0].get("generated_text", "No output generated.")
         elif isinstance(api_response, dict):  # Handle dict response
-            return api_response.get("generated_text", "No output generated.")
+            output = api_response.get("generated_text", "No output generated.")
         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
     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.")