Update app.py

app.py

@@ -384,10 +384,10 @@ def extract_functions(gemini_output):
 
 def validate_and_generate_documentation(api_url, headers, gemini_output, file_contents, functionality_description):
     """
-    Uses the Hugging Face Inference API to generate complete documentation without truncation or prompts.
+    Uses the Hugging Face Inference API to generate documentation in chunks to avoid token limits.
     """
     # Generate the refined prompt for the Qwen model
-    prompt = f"""
+    base_prompt = f"""
     User-specified functionality: '{functionality_description}'
     Functions identified by Gemini:
     {gemini_output}
@@ -419,43 +419,43 @@ def validate_and_generate_documentation(api_url, headers, gemini_output, file_co
           - Example Usage: <Example demonstrating how to use the function>
     """
 
-    # Send the prompt to the Hugging Face API
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 2048}}
-    response = requests.post(api_url, headers=headers, json=payload)
+    # 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
 
-    if response.status_code == 200:
-        api_response = response.json()
-        full_output = ""
+    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)
 
-        # Handle multi-part outputs by looping until all content is retrieved
-        while isinstance(api_response, list) or isinstance(api_response, dict):
+    # 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:
-                break
-
+                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}")
 
-            # If the response indicates more content, request the next chunk
-            if "end_of_text" in output or len(output.strip()) == 0:
-                break
-
-            # Prepare the payload for continuation
-            payload = {"inputs": full_output, "parameters": {"max_new_tokens": 2048}}
-            response = requests.post(api_url, headers=headers, json=payload)
-
-            if response.status_code == 200:
-                api_response = response.json()
-            else:
-                raise ValueError(f"Error during continuation: {response.status_code}, {response.text}")
+    return full_output
 
-        # Clean the output to exclude prompts
-        cleaned_output = clean_output(full_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}")