autodoc app.py

app.py

@@ -0,0 +1,159 @@
+import os
+import sys
+import io
+from dotenv import load_dotenv
+import google.generativeai as genai
+from openai import OpenAI
+import gradio as gr
+
+# Load environment variables
+load_dotenv()
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+
+# Configure APIs
+genai.configure(api_key=GEMINI_API_KEY)
+openai = OpenAI(api_key=OPENAI_API_KEY)
+
+# Model names
+GEMINI_MODEL = "gemini-2.0-flash"  # Gemini model
+OPENAI_MODEL = "gpt-4o-mini"  # OpenAI model (you can use "gpt-3.5-turbo" if needed)
+
+def create_system_prompt():
+    """Create the system prompt for the model."""
+    system_message = "You are an assistant that adds high-quality docstrings and comments to Python code. "
+    system_message += "Your task is to analyze the code and add appropriate documentation following PEP 257 conventions. "
+    system_message += "Add function/class/module docstrings and inline comments where helpful. "
+    system_message += "Don't change the functionality of the code, only add documentation. "
+    system_message += "Use descriptive docstrings that explain parameters, return values, and exceptions raised."
+    return system_message
+
+def create_user_prompt(code):
+    """Create the user prompt with the code to be documented."""
+    user_prompt = "Please add appropriate docstrings and comments to the following Python code. "
+    user_prompt += "Follow PEP 257 style for docstrings. Add comments only where they add value. "
+    user_prompt += "Don't change the functionality of the code.\n\n"
+    user_prompt += code
+    return user_prompt
+
+def document_code_with_gemini(code):
+    """Generate documentation for the provided code using Gemini."""
+    system_prompt = create_system_prompt()
+    user_prompt = create_user_prompt(code)
+    
+    # Use Gemini to generate documentation
+    model = genai.GenerativeModel(GEMINI_MODEL)
+    response = model.generate_content(
+        f"{system_prompt}\n\n{user_prompt}",
+        generation_config=genai.GenerationConfig(
+            temperature=0.7,
+            max_output_tokens=2048,
+        ),
+    )
+    
+    # Clean up the response
+    documented_code = response.text.replace("```python", "").replace("```", "").strip()
+    return documented_code
+
+def document_code_with_openai(code):
+    """Generate documentation for the provided code using OpenAI."""
+    system_prompt = create_system_prompt()
+    user_prompt = create_user_prompt(code)
+    
+    # Use OpenAI to generate documentation
+    response = openai.chat.completions.create(
+        model=OPENAI_MODEL,
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        temperature=0.7,
+        max_tokens=2048,
+    )
+    
+    # Clean up the response
+    documented_code = response.choices[0].message.content.replace("```python", "").replace("```", "").replace("python", "")
+    return documented_code
+
+def document_code(code, api_choice):
+    """Generate documentation using the selected API."""
+    if api_choice == "Gemini":
+        return document_code_with_gemini(code)
+    elif api_choice == "OpenAI":
+        return document_code_with_openai(code)
+    else:
+        return "Invalid API choice."
+
+def execute_python(code):
+    """Execute Python code and return the output."""
+    try:
+        output = io.StringIO()
+        sys.stdout = output
+        exec(code)
+        return output.getvalue()
+    except Exception as e:
+        return f"Error: {str(e)}"
+    finally:
+        sys.stdout = sys.__stdout__
+
+def create_interface():
+    """Create the Gradio interface for the documentation tool."""
+    sample_code = """
+def fibonacci(n):
+    fib_series = []
+    a, b = 0, 1
+    for _ in range(n):
+        fib_series.append(a)
+        a, b = b, a + b
+    return fib_series
+
+n = 10  
+fib_series = fibonacci(n)
+print("Fibonacci series up to", n, "terms:", fib_series)
+"""
+
+    css = """
+    .original {background-color: #2d2d2d; color: #f8f8f2;}
+    .documented {background-color: #1e1e1e; color: #f8f8f2;}
+
+    /* Custom button styles */
+    #add-documentation-btn {background-color: #007bff; color: white; font-weight: bold;} /* Blue for Add Documentation */
+    #run-original-btn {background-color: #28a745; color: white; font-weight: bold;} /* Green for Run Original Code */
+    #run-documented-btn {background-color: #6f42c1; color: white; font-weight: bold;} /* Purple for Run Documented Code */
+
+    /* Ensure all buttons have white bold text */
+    button {color: white; font-weight: bold;}
+    """
+
+    with gr.Blocks(css=css) as ui:
+        gr.Markdown("## Automatic Python Documentation Generator")
+        
+        with gr.Row():
+            input_code = gr.Textbox(label="Original Python code:", value=sample_code, lines=15, elem_classes=["original"])
+            documented_code = gr.Textbox(label="Documented Python code:", lines=15, elem_classes=["documented"])
+        
+        with gr.Row():
+            api_choice = gr.Radio(["Gemini", "OpenAI"], label="Select API", value="Gemini")
+            document_btn = gr.Button("Add Documentation", elem_id="add-documentation-btn")
+        
+        with gr.Row():
+            run_original = gr.Button("Run Original Code", elem_id="run-original-btn")
+            run_documented = gr.Button("Run Documented Code", elem_id="run-documented-btn")
+        
+        with gr.Row():
+            original_output = gr.TextArea(label="Original code output:")
+            documented_output = gr.TextArea(label="Documented code output:")
+        
+        document_btn.click(document_code, inputs=[input_code, api_choice], outputs=[documented_code])
+        run_original.click(execute_python, inputs=[input_code], outputs=[original_output])
+        run_documented.click(execute_python, inputs=[documented_code], outputs=[documented_output])
+        
+    return ui
+
+def main():
+    """Main function to launch the interface."""
+    ui = create_interface()
+    ui.launch(inbrowser=True)
+
+if __name__ == "__main__":
+    main()