Build a bot that checks that docstring descriptions match function source

In this notebook, we are going to build an LLM-powered bot that checks that docstring descriptions match function source. We will use the LlamaBot's StructuredBot and Pydantic to make this happen.

Define the behaviour of our bot

The bot's ideal behaviour will look like this:

1. It will be given function's source code.
2. It will then be asked to return a boolean judgment call:
   1. If the docstring matches the function source, the answer will be "True".
   2. If the docstring doesn't match, the answer will be "False" along with a list of reasons.

Define Pydantic model

The desired behaviour above means we need the following Pydantic model:

from pydantic import BaseModel, Field


class DocstringDescribesFunction(BaseModel):
    docs_match_source: bool = Field(
        default=False,
        description="Whether or not the docstring matches the function source.",
    )
    reasons: list[str] = Field(
        default_factory=list,
        description="Reasons why the docstring doesn't match the function source.",
    )

Define the StructuredBot's behaviour

We will now design the prompt for StructuredBot, particularly focusing in on the system prompt for the StructuredBot.

The system prompt is an opportunity for us to steer the behaviour of StructuredBot. Here, we leave instructions for the bot to follow. Doing so here allows us to ensure that the bot's __call__ method only needs to be concerned with receiving a function's source (as a string).

from llamabot import prompt


@prompt
def docstringbot_sysprompt() -> str:
    """You are an expert at documenting functions.

    You will be given a docstring and a function source.
    Your job is to determine if the docstring matches the function source.

    If it does match, respond with no reasons and respond with "True".

    If it doesn't match,
    respond with a list of reasons why the docstring doesn't match the function source.
    Be specific about the reasons, such as:

    - "The docstring is mismatched with the function. The function does <something>,
      but the docstring says <something_else>."
    - "The docstring is completely missing."
    """

Next up, let's create the StructuredBot. Upon initializing, we provide the system prompt and the Pydantic model that it needs to reference.

from llamabot import StructuredBot

docstringbot = StructuredBot(
    docstringbot_sysprompt(), pydantic_model=DocstringDescribesFunction
)

Test docstringbot on different functions

def fibbonacci(n: int) -> int:
    """Return the nth Fibonacci number.

    Mathematically, the nth Fibonnaci number is defined as
    the sum of the (n-1)th and (n-2)th Fibonacci numbers.

    As such, this is what is returned.

    :param n: The position of the Fibonacci number to return.
    """
    if n <= 0:
        raise ValueError("n must be a positive integer.")
    elif n == 1:
        return 0
    elif n == 2:
        return 1
    else:
        return fibbonacci(n - 1) + fibbonacci(n - 2)

from inspect import getsource

source_code = getsource(fibbonacci)
docstringbot(source_code)

Let's try now an example where the docstring is completely missing.

def fibbonacci(n: int) -> int:
    if n <= 0:
        raise ValueError("n must be a positive integer.")
    elif n == 1:
        return 0
    elif n == 2:
        return 1
    else:
        return fibbonacci(n - 1) + fibbonacci(n - 2)


source_code = getsource(fibbonacci)
docstringbot(source_code)

And now let's try an example where the docstring doesn't match the function source.

def fibbonacci(n: int) -> int:
    """This function bakes a cake of the Fibonacci sequence."""
    if n <= 0:
        raise ValueError("n must be a positive integer.")
    elif n == 1:
        return 0
    elif n == 2:
        return 1
    else:
        return fibbonacci(n - 1) + fibbonacci(n - 2)


source_code = getsource(fibbonacci)
docstringbot(source_code)