You are currently viewing Module 3: AI-Powered Docstring Generation with DeepSeek and Flask

Module 3: AI-Powered Docstring Generation with DeepSeek and Flask

  • Post author:
  • Post category:Python

Welcome to Module 3 of the tutorial series: “Build an AI-powered documentation assistant with Flask & DeepSeek”. In this module, you’ll use DeepSeek’s language model and Flask to generate and improve docstrings.

Prerequisites

Before starting Module 3, ensure you’ve completed the following steps:

Lesson 5: Generating docstrings with DeepSeek LLM

Objective

In this lesson, you’ll use DeepSeek’s language model to generate docstrings for functions and classes. You’ll also format the docstrings to adhere to PEP-257 standards and test their accuracy.

Step 1: Update docstring_generator.py

This file will contain all the logic for generating docstrings using DeepSeek.

File to Update: app/utils/docstring_generator.py

import os
import requests
from dotenv import load_dotenv

load_dotenv()

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")

def generate_docstring(code_snippet, context=None):
    """
    Generate a docstring for a given code snippet using DeepSeek.
    """
    prompt = f"""
    Generate a docstring for the following Python code. Follow PEP-257 standards and include:
    - A one-line summary.
    - A detailed description (if necessary).
    - Args (for functions).
    - Returns (for functions).
    - Raises (if applicable).

    Code:
    {code_snippet}

    Context:
    {context if context else "No additional context provided."}
    """

    headers = {
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",  # Specify the model
        "prompt": prompt,
        "max_tokens": 200,
        "temperature": 0.7
    }
    response = requests.post(
        "https://api.deepseek.com/beta/completions",
        headers=headers,
        json=data
    )
    if response.status_code == 200:
        return response.json()["choices"][0]["text"].strip()
    else:
        # Print the response details for debugging
        print(f"API Response: {response.status_code}, {response.text}")
        raise Exception(f"Failed to generate docstring: {response.status_code}")

The code generates a docstring for a given Python code snippet using the DeepSeek API. It follows these steps:

  1. Loads environment variables – The script imports os, requests, and load_dotenv from dotenv. It calls load_dotenv() to load environment variables from a .env file.
  2. Retrieves the API key – The script fetches the DEEPSEEK_API_KEY from environment variables using os.getenv("DEEPSEEK_API_KEY").
  3. Generates a prompt – The function generate_docstring() constructs a detailed prompt following PEP-257 standards. It includes:
    • A one-line summary.
    • A detailed description if necessary.
    • Function arguments (Args).
    • Return values (Returns).
    • Raised exceptions (Raises).
  4. Sends an API request – The function sends a POST request to the DeepSeek API using the requests library. It:
    • Sets the Authorization header with the API key.
    • Defines request parameters like model, prompt, max_tokens, and temperature.
    • Sends the JSON payload to "https://api.deepseek.com/beta/completions".
    • Uses the deepseek-chat model, which is designed for natural language understanding and code-related tasks, ensuring accurate and well-structured docstring generation.
  5. Handles API responses – If the API responds successfully (status_code == 200), the function extracts and returns the generated docstring.
  6. Handles errors – If the API request fails, the function prints the response details for debugging and raises an exception.

Step 2: Update routes.py

Update:generate_docstring_route()function to contain the following snippet:

def generate_docstring_route():
    """
    Generate a docstring for a given code snippet.
    """
    data = request.json
    code_snippet = data.get("code")
    context = data.get("context", "")

    try:
        docstring = generate_docstring(code_snippet, context)
        return jsonify({"docstring": docstring})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

The generate_docstring_route() function handles a Flask route that generates a docstring for a given Python code snippet. Here’s how it works:

  1. Extracts JSON data from the request – The function retrieves the JSON payload from the incoming HTTP request using request.json.
  2. Gets the code snippet and context – It extracts the code field from the JSON data, which contains the Python code snippet. If the request includes a context field, it retrieves that as well; otherwise, it sets the context to an empty string.
  3. Generates the docstring – The function calls generate_docstring() with the extracted code_snippet and context. This function interacts with the DeepSeek API to generate a docstring.
  4. Returns the generated docstring – If the API call succeeds, the function sends a JSON response containing the generated docstring.
  5. Handles errors – If an error occurs, the function catches the exception, converts it into a JSON error message, and returns an HTTP 500 status code.

This function enables users to submit Python code and receive an automatically generated docstring in response.

Step 3: Test the /generate-docstring Endpoint

  1. Start the Flask app: 
    python run.py
  2. Use curl to call the /generate-docstring endpoint: 
    curl -X POST http://127.0.0.1:5000/generate-docstring \
     -H "Content-Type: application/json" \
     -d '{"code": "def add(a, b): return a + b"}'
  3. Expected output: 
    {
  "docstring": "```python\n\
def add(a, b):\n\
    \"\"\"Return the sum of two numbers.\n\n\
    This function takes two numeric inputs and returns their sum.\n\n\
    Args:\n\
        a (int or float): The first number to be added.\n\
        b (int or float): The second number to be added.\n\n\
    Returns:\n\
        int or float: The sum of the two input numbers.\n\n\
    Raises:\n\
        TypeError: If either `a` or `b` is not a number (int or float).\n\
    \"\"\"\n\
    return a + b\n\
```"
}

     

Lesson 6: Improving existing docstrings with AI

Objective

In this lesson, you’ll use DeepSeek to improve existing docstrings by enhancing clarity, readability, and completeness.

Step 1: Update docstring_generator.py

Add a function to improve existing docstrings.

File to Update: app/utils/docstring_generator.py

def improve_docstring(existing_docstring, context=None):
    """
    Improve an existing docstring using DeepSeek.
    """
    prompt = f"""
    Improve the following docstring for clarity, readability, and completeness. Convert passive voice to active voice where applicable.

    Docstring:
    {existing_docstring}

    Context:
    {context if context else "No additional context provided."}
    """

    headers = {
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "deepseek-chat",  # Specify the model
        "prompt": prompt,
        "max_tokens": 200,
        "temperature": 0.7
    }
    response = requests.post(
        "https://api.deepseek.com/beta/completions",  # Use the beta endpoint
        headers=headers,
        json=data
    )
    if response.status_code == 200:
        return response.json()["choices"][0]["text"].strip()
    else:
        # Print the response details for debugging
        print(f"API Response: {response.status_code}, {response.text}")
        raise Exception(f"Failed to improve docstring: {response.status_code}")

This function enhances an existing docstring using the DeepSeek API. Here’s how it works:

  1. Constructs a prompt – The function improve_docstring() creates a structured prompt that instructs DeepSeek to refine the given docstring. The prompt ensures the improvement focuses on clarity, readability, completeness, and converting passive voice to active voice.
  2. Includes context information – If provided, the function appends additional context to the prompt. If no context is available, it uses a default message stating that no extra information is provided.
  3. Prepares API request – The function sets up request headers, including the Authorization header with the API key and the content type as application/json. It then defines the request payload, specifying:
    • The model (deepseek-chat).
    • The generated prompt.
    • The maximum token limit (max_tokens = 200).
    • The temperature (temperature = 0.7) to balance creativity and coherence.
  4. Sends API request – The function makes a POST request to the DeepSeek API at "https://api.deepseek.com/beta/completions".
  5. Handles API response – If the API responds successfully (status_code == 200), the function extracts the improved docstring from the response and returns it.
  6. Handles errors – If the API request fails, the function prints the response details for debugging and raises an exception with the error status code.

Step 2: Update routes.py

Add a new route to improve existing docstrings.

File to Update: app/routes.py

from .utils.docstring_generator import generate_docstring, improve_docstring

@main_bp.route("/improve-docstring", methods=["POST"])
def improve_docstring_route():
    """
    Improve an existing docstring.
    """
    data = request.json
    existing_docstring = data.get("docstring")
    context = data.get("context", "")

    try:
        improved_docstring = improve_docstring(existing_docstring, context)
        return jsonify({"improved_docstring": improved_docstring})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

This code defines an API route that enhances an existing docstring using the improve_docstring function. Here’s how it works:

  1. Defines an API endpoint – The @main_bp.route("/improve-docstring", methods=["POST"]) decorator registers the improve_docstring_route() function as a POST endpoint at /improve-docstring.
  2. Extracts input data – The function retrieves JSON data from the request, extracting the docstring and an optional context.
  3. Calls improve_docstring – It passes the existing docstring and context to the improve_docstring function, which refines the content for clarity, readability, and completeness.
  4. Returns the improved docstring – If the function successfully enhances the docstring, it responds with a JSON object containing the updated version.
  5. Handles errors – If an exception occurs, the function catches it, formats the error message as JSON, and returns an HTTP 500 response.

Step 3: Test the /improve-docstring Endpoint

  1. Start the Flask app: 
    python run.py
  2. Use curl to call the /improve-docstring endpoint: 
    curl -X POST http://127.0.0.1:5000/improve-docstring \
     -H "Content-Type: application/json" \
     -d '{"docstring": "Adds two numbers together."}'
  3. Expected output: 
    {
  "improved_docstring": "1. Add a clear and concise description of the function's purpose.\n"
    "2. Specify the input parameters and their types.\n"
    "3. Specify the return value and its type.\n"
    "4. Include any relevant examples or usage.\n"
    "5. Convert passive voice to active voice where applicable.\n\n"
    "Improved docstring:\n"
    "\"\"\"\n"
    "Adds two numbers together.\n\n"
    "Parameters:\n"
    "a (int or float): The first number to add.\n"
    "b (int or float): The second number to add.\n\n"
    "Returns:\n"
    "int or float: The sum of the two numbers.\n\n"
    "Examples:\n"
    ">>> add(2, 3)\n"
    "5\n"
    ">>> add(2.5, 3.7)\n"
    "6.2\n"
    "\"\"\""
}

What you’ve achieved

  • You generated docstrings for functions and classes using DeepSeek.
  • You improved existing docstrings for clarity, readability, and completeness.
  • You built Flask API endpoints to automate docstring generation and improvement.

 

Full code for module 3

You can find the complete code for this tutorial in the GitHub repository.

Next Steps

  • Proceed to Module 4: Learn how to implement an interactive AI documentation search system.
  • Experiment: Try generating and improving docstrings for more complex code snippets.
  • Join the Community: Share your progress and get feedback from other learners!

Facebook Comments

Tags: ,