Repository: airtai/docstring-gen (language: python)

Issue:
Create a CLI command to generate docstring
- [x] Write a function to read and write notebook cells using AST
- [x] Add support for python files
- [x] Write a function to use open AI model to generate docstring for functions and classes
- [x] Handle rate limits
- [x] Wrap the above two into a CLI command 'docstring gen'


Discussion / hints:

Fix patch:
diff --git a/docstring_gen/_cli.py b/docstring_gen/_cli.py
new file mode 100644
index 0000000..1170324
--- /dev/null
+++ b/docstring_gen/_cli.py
@@ -0,0 +1,76 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/CLI.ipynb.
+
+# %% auto 0
+__all__ = ['gen']
+
+# %% ../nbs/CLI.ipynb 2
+from typing import *
+from pathlib import Path
+import importlib
+
+import typer
+
+from .docstring_generator import add_docstring_to_source
+
+# %% ../nbs/CLI.ipynb 5
+_app = typer.Typer()
+
+
+@_app.command(
+    help="This command reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings to classes and methods that do not have them.",
+)
+def gen(
+    path: str = typer.Argument(
+        ".",
+        help="The path to the Jupyter notebook or Python file, or a directory containing these files",
+    ),
+    include_auto_gen_txt: bool = typer.Option(
+        True,
+        help="If set to True, a note indicating that the docstring was autogenerated by docstring-gen library will be added to the end.",
+    ),
+    recreate_auto_gen_docs: bool = typer.Option(
+        False,
+        "--force-recreate-auto-generated",
+        "-f",
+        help="If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.",
+    ),
+    model: str = typer.Option(
+        "code-davinci-002",
+        help="The name of the Codex model that will be used to generate docstrings.",
+    ),
+    temperature: float = typer.Option(
+        0.2,
+        help="Setting the temperature close to zero produces better results, whereas higher temperatures produce more complex, and sometimes irrelevant docstrings.",
+        min=0.0,
+        max=1.0,
+    ),
+    max_tokens: int = typer.Option(
+        250,
+        help="The maximum number of tokens to be used when generating a docstring for a function or class. Please note that a higher number will deplete your token quota faster.",
+    ),
+    top_p: float = typer.Option(
+        1.0,
+        help="You can also specify a top-P value from 0-1 to achieve similar results to changing the temperature. According to the Open AI documentation, it is generally recommended to change either this or the temperature but not both.",
+        min=0.0,
+        max=1.0,
+    ),
+    n: int = typer.Option(
+        3,
+        help="The number of docstrings to be generated for each function or class, with the best one being added to the source code. Please note that a higher number will deplete your token quota faster.",
+    ),
+) -> None:
+
+    try:
+        add_docstring_to_source(
+            path=path,
+            include_auto_gen_txt=include_auto_gen_txt,
+            recreate_auto_gen_docs=recreate_auto_gen_docs,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            n=n,
+        )
+    except Exception as e:
+        typer.secho(e, err=True, fg=typer.colors.RED)
+        raise typer.Exit(1)
diff --git a/docstring_gen/_modidx.py b/docstring_gen/_modidx.py
index f3d89c9..0f53379 100644
--- a/docstring_gen/_modidx.py
+++ b/docstring_gen/_modidx.py
@@ -5,21 +5,35 @@
                 'doc_host': 'https://airtai.github.io',
                 'git_url': 'https://github.com/airtai/docstring-gen',
                 'lib_path': 'docstring_gen'},
-  'syms': { 'docstring_gen.docstring_generator': { 'docstring_gen.docstring_generator._add_docstring': ( 'docstring_generator.html#_add_docstring',
-                                                                                                         'docstring_gen/docstring_generator.py'),
-                                                   'docstring_gen.docstring_generator._calculate_end_lineno': ( 'docstring_generator.html#_calculate_end_lineno',
-                                                                                                                'docstring_gen/docstring_generator.py'),
+  'syms': { 'docstring_gen.docstring_generator': { 'docstring_gen.docstring_generator._add_auto_gen_txt': ( 'docstring_generator.html#_add_auto_gen_txt',
+                                                                                                            'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._add_docstring_to_nb': ( 'docstring_generator.html#_add_docstring_to_nb',
+                                                                                                               'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._add_docstring_to_py': ( 'docstring_generator.html#_add_docstring_to_py',
+                                                                                                               'docstring_gen/docstring_generator.py'),
                                                    'docstring_gen.docstring_generator._check_and_add_docstrings_to_source': ( 'docstring_generator.html#_check_and_add_docstrings_to_source',
                                                                                                                               'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._completions_with_backoff': ( 'docstring_generator.html#_completions_with_backoff',
+                                                                                                                    'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._fix_docstring_indent': ( 'docstring_generator.html#_fix_docstring_indent',
+                                                                                                                'docstring_gen/docstring_generator.py'),
                                                    'docstring_gen.docstring_generator._generate_docstring_using_codex': ( 'docstring_generator.html#_generate_docstring_using_codex',
                                                                                                                           'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._get_best_docstring': ( 'docstring_generator.html#_get_best_docstring',
+                                                                                                              'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._get_classes_and_functions': ( 'docstring_generator.html#_get_classes_and_functions',
+                                                                                                                     'docstring_gen/docstring_generator.py'),
                                                    'docstring_gen.docstring_generator._get_code_from_source': ( 'docstring_generator.html#_get_code_from_source',
                                                                                                                 'docstring_gen/docstring_generator.py'),
-                                                   'docstring_gen.docstring_generator._get_start_line_for_class_or_func': ( 'docstring_generator.html#_get_start_line_for_class_or_func',
-                                                                                                                            'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._get_files': ( 'docstring_generator.html#_get_files',
+                                                                                                     'docstring_gen/docstring_generator.py'),
                                                    'docstring_gen.docstring_generator._inject_docstring_to_source': ( 'docstring_generator.html#_inject_docstring_to_source',
                                                                                                                       'docstring_gen/docstring_generator.py'),
-                                                   'docstring_gen.docstring_generator._line_has_class_or_method': ( 'docstring_generator.html#_line_has_class_or_method',
-                                                                                                                    'docstring_gen/docstring_generator.py'),
-                                                   'docstring_gen.docstring_generator.add_docstring_to_notebook': ( 'docstring_generator.html#add_docstring_to_notebook',
-                                                                                                                    'docstring_gen/docstring_generator.py')}}}
+                                                   'docstring_gen.docstring_generator._remove_auto_generated_docstring': ( 'docstring_generator.html#_remove_auto_generated_docstring',
+                                                                                                                           'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._retry_with_exponential_backoff': ( 'docstring_generator.html#_retry_with_exponential_backoff',
+                                                                                                                          'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator._visit_functions': ( 'docstring_generator.html#_visit_functions',
+                                                                                                           'docstring_gen/docstring_generator.py'),
+                                                   'docstring_gen.docstring_generator.add_docstring_to_source': ( 'docstring_generator.html#add_docstring_to_source',
+                                                                                                                  'docstring_gen/docstring_generator.py')}}}
diff --git a/docstring_gen/docstring_generator.py b/docstring_gen/docstring_generator.py
index f2e8dca..5563b04 100644
--- a/docstring_gen/docstring_generator.py
+++ b/docstring_gen/docstring_generator.py
@@ -1,232 +1,646 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Docstring_Generator.ipynb.
 
 # %% auto 0
-__all__ = ['add_docstring_to_notebook']
+__all__ = ['DEFAULT_PROMPT_TEMPLATE', 'DOCSTRING_ERR', 'AUTO_GEN_PERFIX', 'AUTO_GEN_BODY', 'AUTO_GEN_SUFFIX', 'AUTO_GEN_TXT',
+           'add_docstring_to_source']
 
 # %% ../nbs/Docstring_Generator.ipynb 2
+import time
+import random
 import ast
-import tokenize
+import textwrap
+import os
+import re
 from typing import *
 from pathlib import Path
-from io import BytesIO
 
 import nbformat
+import openai
+import typer
 
-# %% ../nbs/Docstring_Generator.ipynb 5
-def _generate_docstring_using_codex(code: str) -> str:
-    return """Sample docstring
+from mypy_extensions import NamedArg
+
+# %% ../nbs/Docstring_Generator.ipynb 4
+def _visit_functions(
+    tree: ast.AST,
+    *,
+    source: str,
+    start_lineno: int = 1,
+    end_lineno: int = -1,
+    callback: Callable[
+        [ast.AST, str, int, int, NamedArg(List[Tuple[int, int, int, int]], "retval")],
+        Any,
+    ],
+    **kwargs: Any
+) -> None:
+    """Walk the abstract syntax tree and call the callback function for every node found
 
     Args:
-        s: sample args
+        tree: The python AST
+        source: The source code
+        start_lineno: The start line number
+        end_lineno: The end line number
+        callback: The callback function
+        kwargs: The keyword arguments
 
-    Returns:
-        sample return
-"""
+    !!! note
 
-# %% ../nbs/Docstring_Generator.ipynb 6
-def _inject_docstring_to_source(
-    source: str, docstring: str, lineno: int, node_col_offset: int
-) -> str:
-    """Inject a docstring into the source code at a specified line number.
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    if end_lineno == -1:
+        end_lineno = len(source.split("\n"))
+
+    if isinstance(tree, (ast.AsyncFunctionDef, ast.FunctionDef, ast.ClassDef)):
+        if ast.get_docstring(tree) is None:
+            callback(tree, source, start_lineno, end_lineno, **kwargs)
+
+    if hasattr(tree, "body"):
+        for i, n in enumerate(tree.body):
+            node_start_lineno = tree.body[i].lineno
+            node_end_lineno = (
+                tree.body[i + 1].lineno - 1 if i < (len(tree.body) - 1) else end_lineno
+            )
+            _visit_functions(
+                n,
+                source=source,
+                start_lineno=node_start_lineno,
+                end_lineno=node_end_lineno,
+                callback=callback,
+                **kwargs
+            )
+
+# %% ../nbs/Docstring_Generator.ipynb 8
+def _get_classes_and_functions(source: str) -> List[Tuple[int, int, int, int]]:
+    """Get the classes and functions in a source file.
 
     Args:
-        source: the source code
-        docstring: the docstring to be added
-        lineno: the line number at which the docstring will be inserted
-        node_col_offset: the number of spaces to indent the docstring
+        source: The source code of the file.
+        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.
+
 
     Returns:
-        The updated source code with the docstring injected
+        A list of tuples of the form (start_lineno, end_lineno, start_col_offset, end_col_offset)
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    lines = source.split("\n")
-    indented_docstring = "\n".join(
-        [
-            line
-            if i == 0 or i == len(docstring.split("\n")) - 1
-            else f"{' ' * node_col_offset}{line}"
-            for i, line in enumerate(docstring.split("\n"))
-        ]
+    retval: List[Tuple[int, int, int, int]] = []
+    tree = ast.parse(source)
+
+    def callback(tree, source, start_lineno, end_lineno, *, retval):
+        """Callback function for the ast.walk function.
+
+        Args:
+            tree: The tree to walk
+            source: The source code
+            start_lineno: The starting line number
+            end_lineno: The ending line number
+
+        Keyword Args:
+            retval: The return value
+
+        Returns:
+            The return value
+
+        !!! note
+
+            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+        """
+        retval.append(
+            (start_lineno, tree.body[0].lineno - 1, end_lineno, tree.body[0].col_offset)
+        )
+
+    _visit_functions(
+        tree,
+        source=source,
+        callback=callback,
+        retval=retval,
     )
-    indent = node_col_offset + 4
-    lines.insert(lineno, f'{" " * indent}"""{indented_docstring}{" " * indent}"""')
-    return "\n".join(lines)
 
-# %% ../nbs/Docstring_Generator.ipynb 8
+    return retval
+
+# %% ../nbs/Docstring_Generator.ipynb 10
 def _get_code_from_source(source: str, start_line_no: int, end_line_no: int) -> str:
-    """Get a block of lines from a given source string.
+    """Get code from source
 
     Args:
-        source: The source string.
-        start_line_no: The line number of the start of the block of lines.
-        end_line_no: The line number of the end of the block of lines.
+        source: The source code of the file.
+        start_line_no: Start line number
+        end_line_no: End line number
 
     Returns:
-        The extracted block of lines from the source
+        The extracted code
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
     source_lines = source.split("\n")
     extracted_lines = source_lines[start_line_no - 1 : end_line_no]
     return "\n".join(extracted_lines)
 
-# %% ../nbs/Docstring_Generator.ipynb 10
-def _calculate_end_lineno(source: str, start_line_no: int) -> int:
-    """Calculate the end line number of a function in a Python source code.
+# %% ../nbs/Docstring_Generator.ipynb 12
+# Reference: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_handle_rate_limits.ipynb
+
+
+def _retry_with_exponential_backoff(
+    initial_delay: float = 1,
+    exponential_base: float = 2,
+    jitter: bool = True,
+    max_retries: int = 10,
+    max_wait: float = 60,
+    errors: tuple = (
+        openai.error.RateLimitError,
+        openai.error.ServiceUnavailableError,
+        openai.error.APIError,
+    ),
+) -> Callable:
+    """Retry a function with exponential backoff."""
+
+    def decorator(func):
+        """Decorator to retry a function if it fails.
+
+        Args:
+            func: The function to be decorated
+            max_retries: The maximum number of retries
+            initial_delay: The initial delay
+            exponential_base: The exponential base
+            max_wait: The maximum wait
+            jitter: The jitter
+
+        Returns:
+            The decorated function
+
+        Raises:
+            Exception: If the maximum number of retries is exceeded
+
+        !!! note
+
+            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+        """
+
+        def wrapper(*args, **kwargs):
+            num_retries = 0
+            delay = initial_delay
+
+            while True:
+                try:
+                    return func(*args, **kwargs)
+
+                except errors as e:
+                    num_retries += 1
+                    if num_retries > max_retries:
+                        raise Exception(
+                            f"Maximum number of retries ({max_retries}) exceeded."
+                        )
+                    delay = min(
+                        delay
+                        * exponential_base
+                        * (1 + jitter * random.random()),  # nosec
+                        max_wait,
+                    )
+                    typer.secho(
+                        f"Note: OpenAI's API rate limit reached. Command will automatically retry in {int(delay)} seconds. For more information visit: https://help.openai.com/en/articles/5955598-is-api-usage-subject-to-any-rate-limits",
+                        fg=typer.colors.BLUE,
+                    )
+                    time.sleep(delay)
+
+                except Exception as e:
+                    raise e
+
+        return wrapper
+
+    return decorator
+
+
+@_retry_with_exponential_backoff()
+def _completions_with_backoff(*args, **kwargs):
+    return openai.Completion.create(*args, **kwargs)
+
+# %% ../nbs/Docstring_Generator.ipynb 13
+DEFAULT_PROMPT_TEMPLATE = '''
+# Python 3.7
+
+#docstring_gen: Original code:
+
+def add_strings(s1: Optional[str] = None, s2: Optional[str] = None) -> str:
+    if s1 is None or s2 is None:
+        raise ValueError(Both s1 and s2 must be provided and must be of type string")
+    return s1 + s2
+    
+#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.
+
+"""Add two strings
+
+Args:
+    s1: First string
+    s2: Second string
+    
+Returns:
+    The added string
+    
+Raises:
+    ValueError: If s1 or s2 is None
+"""
+
+#docstring_gen: Original code:
+
+class Person:
+    def __init__(self, name, surname, age):
+        self.name = name
+        self.surname = surname
+        self.age = age
+
+    def info(self, additional=""):
+        print('My name is :' + self.name + additional)
+        
+#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.
+
+
+"""A class to represent a person.
+
+Attributes:
+    name : first name of the person
+    surname : family name of the person
+    age : age of the person
+"""
+
+#docstring_gen: Original code:
+
+{code}
+
+#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.
+
+"""
+'''
+
+# %% ../nbs/Docstring_Generator.ipynb 18
+def _get_best_docstring(docstrings: List[str]) -> Optional[str]:
+    """Get the best docstring from a list of docstrings.
 
     Args:
-        source: The source code string.
-        start_line_no: The line number of the start of the function.
+        docstrings: List of docstrings
 
     Returns:
-        The end line number of the function.
+        The best docstring
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    lines = source.split("\n")[start_line_no - 1 :]
-    first_indent = len(lines[0]) - len(lines[0].lstrip())
-    end_line_in_source = 0
-
-    for i, line in enumerate(lines[1:]):
-        if len(line) - len(line.lstrip()) == first_indent and line.strip() != "":
-            end_line_in_source = i
-            break
-
-    ret_val = (
-        len(source.split("\n"))
-        if end_line_in_source == 0
-        else end_line_in_source + start_line_no
-    )
-    return ret_val - 1
+    docstrings = [d for d in docstrings if "Args:" in d or "Attributes:" in d]
+    docstrings = [d for d in docstrings if "~~~~" not in d]
+    return docstrings[0] if len(docstrings) > 0 else None
+
+# %% ../nbs/Docstring_Generator.ipynb 21
+DOCSTRING_ERR = """!!! note
+    
+    Failed to generate docs
+"""
 
-# %% ../nbs/Docstring_Generator.ipynb 12
-def _line_has_class_or_method(source: str, lineno: int) -> bool:
-    """Check if a line in the source code contains a class or method definition.
+
+def _generate_docstring_using_codex(
+    code: str, prompt_template: str, **kwargs: Union[int, float, str, List[str]]
+) -> str:
+    """Generate a docstring using codex.
 
     Args:
-        source: The source code as a string.
-        lineno: The line number to check.
+        code: The code for which to generate a docstring.
+        prompt_template: The template for the prompt.
+        **kwargs: The keyword arguments.
 
     Returns:
-        True if the line contains a class or method definition, False otherwise.
+        The generated docstring.
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    line = "".join(source.split("\n")[lineno - 1])
-    tokens = list(tokenize.tokenize(BytesIO(line.encode("utf-8")).readline))
-    return tokens[1].type == tokenize.NAME and tokens[1].string in {
-        "class",
-        "def",
-        "async",
-    }
+    prompt = prompt_template.format(code=code)
+    response = _completions_with_backoff(prompt=prompt, **kwargs)
 
+    best_docstring = _get_best_docstring([i.text for i in response.choices])
+    retval = best_docstring if best_docstring is not None else DOCSTRING_ERR
+    return retval
 
-def _get_start_line_for_class_or_func(source: str, lineno: int) -> int:
-    """Get the line number of the first line containing a class or function definition.
+# %% ../nbs/Docstring_Generator.ipynb 23
+AUTO_GEN_PERFIX = """
+!!! note
+
+"""
+AUTO_GEN_BODY = "The above docstring is autogenerated by docstring-gen library"
+AUTO_GEN_SUFFIX = "(https://github.com/airtai/docstring-gen)"
+
+AUTO_GEN_TXT = AUTO_GEN_PERFIX + " " * 4 + AUTO_GEN_BODY + " " + AUTO_GEN_SUFFIX
+
+# %% ../nbs/Docstring_Generator.ipynb 24
+def _add_auto_gen_txt(docstring: str) -> str:
+    """Add the autogenerated by docstring-gen library text to the end of a docstring
 
     Args:
-        source: The source code as a string.
-        lineno: The line number to start from.
+        docstring: The docstring to which the text will be added
 
     Returns:
-        The line number of the first line containing a class or function definition,
-        or the original line number if no such line is found.
+        The docstring with the added text
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    if _line_has_class_or_method(source, lineno):
-        return lineno
-
-    original_lineno = lineno
-    total_lines = source.split("\n")
-    for i in total_lines:
-        lineno += 1
-        if lineno > len(total_lines):
-            break
-        if _line_has_class_or_method(source, lineno):
-            return lineno
-    return original_lineno
-
-# %% ../nbs/Docstring_Generator.ipynb 15
-def _add_docstring(
+    return docstring + AUTO_GEN_TXT + "\n"
+
+# %% ../nbs/Docstring_Generator.ipynb 26
+def _fix_docstring_indent(
+    docstring: str, col_offset: int, *, include_auto_gen_txt: bool
+) -> str:
+    """Fix the indentation of a docstring.
+
+    Args:
+        docstring: The docstring to fix.
+        col_offset: The column offset to use.
+        include_auto_gen_txt: If True, include the auto-generated text in the docstring.
+
+    Returns:
+        The fixed docstring.
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    lines = docstring.split("\n")
+    rest = textwrap.dedent("\n".join(lines[1:]))
+    retval = lines[0] + "\n" + rest
+
+    if include_auto_gen_txt:
+        retval = _add_auto_gen_txt(retval)
+
+    retval = '"""' + retval + '"""'
+    retval = textwrap.indent(retval, prefix=" " * col_offset)
+    return retval
+
+# %% ../nbs/Docstring_Generator.ipynb 30
+def _inject_docstring_to_source(
     source: str,
-    node: Union[ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef],
-    line_offset: int,
-) -> Tuple[str, int]:
-    """Add a docstring to the given node and update the source code.
+    indented_docstrings: List[str],
+    linenos: List[Tuple[int, int, int, int]],
+) -> str:
+    """Injects docstrings into source code
 
     Args:
-        source: the source code from the notebook cell
-        node: the AST node representing a class definition, function definition,
-            or async function definition
-        line_offset: the number of lines added before the current
-            node in the source
+        source: The source code
+        indented_docstrings: The docstrings to be injected
+        linenos: The line numbers where the docstrings should be injected
 
     Returns:
-        A tuple containing the updated source code and the new line number offset
+        The source code with the injected docstrings
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    line_no = node.lineno + line_offset
+    source_lines = source.split("\n")
+    line_offset = 0
+    for docstring, (_, docstring_line_no, _, _) in zip(indented_docstrings, linenos):
+        docstring_lines = docstring.split("\n")
+        source_lines = (
+            source_lines[: docstring_line_no + line_offset]
+            + docstring_lines
+            + source_lines[docstring_line_no + line_offset :]
+        )
+        line_offset += len(docstring_lines)
 
-    # Fix for ast's node.lineno giving line number of decorator
-    # instead of function/class definition in Python 3.7
-    line_no = _get_start_line_for_class_or_func(source, line_no)
+    return "\n".join(source_lines)
 
-    if hasattr(node, "end_lineno") and node.end_lineno is not None:
-        end_line_no = node.end_lineno + line_offset
-    else:
-        end_line_no = _calculate_end_lineno(source, line_no)
+# %% ../nbs/Docstring_Generator.ipynb 32
+def _remove_auto_generated_docstring(source: str) -> str:
+    """Remove the autogenerated docstrings from the source code.
 
-    code = _get_code_from_source(source, line_no, end_line_no)
-    docstring = _generate_docstring_using_codex(code)
+    Args:
+        source: The source code
+
+    Returns:
+        The source code without the auto generated docstring
+
+    !!! note
 
-    source = _inject_docstring_to_source(source, docstring, line_no, node.col_offset)
-    line_offset += len(docstring.split("\n"))
-    return source, line_offset
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    placeholder = "{DOCSTRING_PLACEHOLDER}"
+    retval = re.sub(
+        f'"""((?!""").)*?({AUTO_GEN_BODY}).*?"""', placeholder, source, flags=re.DOTALL
+    )
+    retval = "\n".join([l for l in retval.split("\n") if l.strip() != placeholder])
+    return retval
 
-# %% ../nbs/Docstring_Generator.ipynb 17
-def _check_and_add_docstrings_to_source(source: str) -> str:
-    """Check for missing docstrings in the source code and add them if necessary.
+# %% ../nbs/Docstring_Generator.ipynb 34
+def _check_and_add_docstrings_to_source(
+    source: str,
+    include_auto_gen_txt: bool,
+    recreate_auto_gen_docs: bool,
+    **kwargs: Union[int, float, str, List[str]]
+) -> str:
+    """Check and add docstrings to classes and functions that don't have one.
 
     Args:
-        source: the source code from the notebook cell
+        source: Source code
+        include_auto_gen_txt: Include auto gen text
+        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.
+        kwargs: Keyword arguments
 
     Returns:
-        The updated source code with added docstrings
+        The source code with docstrings
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
+    if recreate_auto_gen_docs:
+        source = _remove_auto_generated_docstring(source)
+
+    linenos = _get_classes_and_functions(source)
+    if len(linenos) != 0:
+        classes_and_functions = [
+            _get_code_from_source(source, start_line_no, end_line_no)
+            for start_line_no, docstring_line_no, end_line_no, node_offset in linenos
+        ]
 
-    tree = ast.parse(source)
-    line_offset = 0
+        docstrings = [
+            _generate_docstring_using_codex(i, DEFAULT_PROMPT_TEMPLATE, **kwargs)
+            for i in classes_and_functions
+        ]
+        offsets = [node_offset for i, _, _, node_offset in linenos]
 
-    for node in tree.body:
-        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
-            continue
-        if ast.get_docstring(node) is not None:
-            continue
-
-        # A class or a function without docstring
-        source, line_offset = _add_docstring(source, node, line_offset)
-        if not isinstance(node, ast.ClassDef):
-            continue
-        # Is a class and we need to check the functions inside
-        # 29 - 36 make it as a recursive function
-        for f in node.body:
-            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                continue
-            if ast.get_docstring(f) is not None:
-                continue
-
-            # should be a function inside the class for which there is no docstring
-            source, line_offset = _add_docstring(source, f, line_offset)
+        indented_docstrings = [
+            _fix_docstring_indent(
+                docstring, offset, include_auto_gen_txt=include_auto_gen_txt
+            )
+            for docstring, offset in zip(docstrings, offsets)
+        ]
+        source = _inject_docstring_to_source(source, indented_docstrings, linenos)
 
     return source
 
-# %% ../nbs/Docstring_Generator.ipynb 19
-def add_docstring_to_notebook(nb_path: Union[str, Path], version: int = 4):
-    """Add docstrings to the source
-
-    This function reads through a Jupyter notebook cell by cell and
-    adds docstrings for classes and methods that do not have them.
+# %% ../nbs/Docstring_Generator.ipynb 37
+def _get_files(p: Path) -> List[Path]:
+    """Get Jupyter notebooks and Python files path in the directory.
 
     Args:
-        nb_path: The notebook file path
-        version: The version of the Jupyter notebook format
+        p: Path to the directory
+
+    Returns:
+        A list of paths to the files in the directory
+
+    Raises:
+        ValueError: If the directory does not contain any Python files or notebooks
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
     """
-    nb_path = Path(nb_path)
-    nb = nbformat.read(nb_path, as_version=version)
+    exts = [".ipynb", ".py"]
+    files = [
+        f
+        for f in p.rglob("*")
+        if f.suffix in exts
+        and not any(p.startswith(".") for p in f.parts)
+        and not f.name.startswith("_")
+    ]
+
+    if len(files) == 0:
+        raise ValueError(
+            f"The directory {p.resolve()} does not contain any Python files or notebooks"
+        )
+
+    return files
+
+# %% ../nbs/Docstring_Generator.ipynb 40
+def _add_docstring_to_nb(
+    file: Path,
+    include_auto_gen_txt: bool,
+    recreate_auto_gen_docs: bool,
+    version: int = 4,
+    **kwargs: Union[int, float, str, List[str]]
+) -> None:
+    """Add docstrings to a Jupyter notebook.
+
+    Args:
+        file (Path): Path to the notebook file.
+        include_auto_gen_txt (bool): If True, include a text indicating that the docstring was auto-generated.
+        recreate_auto_gen_docs (bool): If True, recreate the docstrings even if they already exist.
+        version (int): The version of the notebook.
+        **kwargs: Additional arguments to be passed to the function that generates the docstrings.
 
-    for cell in nb.cells:
+    Returns:
+        None
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    _f = nbformat.read(file, as_version=version)
+    for cell in _f.cells:
         if cell.cell_type == "code":
-            cell["source"] = _check_and_add_docstrings_to_source(cell["source"])
+            cell["source"] = _check_and_add_docstrings_to_source(
+                cell["source"], include_auto_gen_txt, recreate_auto_gen_docs, **kwargs
+            )
+    nbformat.write(_f, file)
 
-    nbformat.write(nb, nb_path)
+
+def _add_docstring_to_py(
+    file: Path,
+    include_auto_gen_txt: bool,
+    recreate_auto_gen_docs: bool,
+    **kwargs: Union[int, float, str, List[str]]
+) -> None:
+    """Add docstrings to a python file.
+
+    Args:
+        file: The path to the python file.
+        include_auto_gen_txt: Whether to include a text indicating that the docstring is auto-generated.
+        recreate_auto_gen_docs: Whether to recreate the docstrings even if they already exist.
+        kwargs: Additional keyword arguments.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If file is not a path to a python file.
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    with file.open("r") as f:
+        source = f.read()
+    source = _check_and_add_docstrings_to_source(
+        source, include_auto_gen_txt, recreate_auto_gen_docs, **kwargs
+    )
+    with file.open("w") as f:
+        f.write(source)
+
+
+def add_docstring_to_source(
+    path: Union[str, Path],
+    include_auto_gen_txt: bool = True,
+    recreate_auto_gen_docs: bool = False,
+    model: str = "code-davinci-002",
+    temperature: float = 0.2,
+    max_tokens: int = 250,
+    top_p: float = 1.0,
+    n: int = 3,
+) -> None:
+    """Reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings to classes and methods that do not have them.
+
+    Args:
+        path: The path to the Jupyter notebook or Python file, or a directory containing these files.
+        include_auto_gen_txt: If set to True, a note indicating that the docstring was autogenerated by docstring-gen library will be added to the end.
+        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.
+        model: The name of the Codex model that will be used to generate docstrings.
+        temperature: Setting the temperature close to zero produces better results, whereas higher temperatures produce more complex, and sometimes irrelevant docstrings.
+        max_tokens: The maximum number of tokens to be used when generating a docstring for a function or class. Please note that a higher number will deplete your token quota faster.
+        top_p: You can also specify a top-P value from 0-1 to achieve similar results to changing the temperature. According to the Open AI documentation, it is generally recommended to change either this or the temperature but not both.
+        n: The number of docstrings to be generated for each function or class, with the best one being added to the source code. Please note that a higher number will deplete your token quota faster.
+
+    Returns:
+        None
+
+    !!! note
+
+        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
+    """
+    path = Path(path)
+    files = _get_files(path) if path.is_dir() else [path]
+    frequency_penalty = 0.0
+    presence_penalty = 0.0
+    stop = ["#", '"""']
+
+    for file in files:
+        if file.suffix == ".ipynb":
+            _add_docstring_to_nb(
+                file=file,
+                include_auto_gen_txt=include_auto_gen_txt,
+                recreate_auto_gen_docs=recreate_auto_gen_docs,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                top_p=top_p,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+                stop=stop,
+                n=n,
+            )
+        else:
+            _add_docstring_to_py(
+                file=file,
+                include_auto_gen_txt=include_auto_gen_txt,
+                recreate_auto_gen_docs=recreate_auto_gen_docs,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                top_p=top_p,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+                stop=stop,
+                n=n,
+            )
diff --git a/fixtures/Test_Data.ipynb b/fixtures/Test_Data.ipynb
index 01b9fa3..c82a104 100644
--- a/fixtures/Test_Data.ipynb
+++ b/fixtures/Test_Data.ipynb
@@ -46,6 +46,33 @@
     "# Title"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63b6e111",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_source = '''\n",
+    "def decorator1(func):\n",
+    "    \"\"\"Decorator function that takes a function as an argument and returns a function.\"\"\"\n",
+    "    pass\n",
+    "    \n",
+    "def decorator2(func):\n",
+    "    \"\"\"Sample docstring\n",
+    "\n",
+    "    Args:\n",
+    "        arg 1: arg 1 description\n",
+    "        arg 2: arg 2 description\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    pass\n",
+    "'''\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/mypy.ini b/mypy.ini
index 4c872a5..5e542d9 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -13,7 +13,7 @@ warn_return_any = True
 show_error_codes = True 
 warn_unused_ignores = True
 
-#disallow_incomplete_defs = True
-#disallow_untyped_decorators = True
+disallow_incomplete_defs = True
+disallow_untyped_decorators = True
 disallow_any_unimported = True
 
diff --git a/nbs/CLI.ipynb b/nbs/CLI.ipynb
new file mode 100644
index 0000000..d26a049
--- /dev/null
+++ b/nbs/CLI.ipynb
@@ -0,0 +1,320 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fbb8ee5c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | default_exp _cli"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f33fd675",
+   "metadata": {},
+   "source": [
+    "# CLI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b56c30cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "from typing import *\n",
+    "from pathlib import Path\n",
+    "import importlib\n",
+    "\n",
+    "import typer\n",
+    "\n",
+    "from docstring_gen.docstring_generator import add_docstring_to_source\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0f96b807",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typer.testing import CliRunner"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "954a97ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "runner = CliRunner()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c39e7c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "_app = typer.Typer()\n",
+    "\n",
+    "\n",
+    "@_app.command(\n",
+    "    help=\"This command reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings to classes and methods that do not have them.\",\n",
+    ")\n",
+    "def gen(\n",
+    "    path: str = typer.Argument(\n",
+    "        \".\",\n",
+    "        help=\"The path to the Jupyter notebook or Python file, or a directory containing these files\",\n",
+    "    ),\n",
+    "    include_auto_gen_txt: bool = typer.Option(\n",
+    "        True,\n",
+    "        help=\"If set to True, a note indicating that the docstring was autogenerated by docstring-gen library will be added to the end.\",\n",
+    "    ),\n",
+    "    recreate_auto_gen_docs: bool = typer.Option(\n",
+    "        False,\n",
+    "        \"--force-recreate-auto-generated\",\n",
+    "        \"-f\",\n",
+    "        help=\"If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\",\n",
+    "    ),\n",
+    "    model: str = typer.Option(\n",
+    "        \"code-davinci-002\",\n",
+    "        help=\"The name of the Codex model that will be used to generate docstrings.\",\n",
+    "    ),\n",
+    "    temperature: float = typer.Option(\n",
+    "        0.2,\n",
+    "        help=\"Setting the temperature close to zero produces better results, whereas higher temperatures produce more complex, and sometimes irrelevant docstrings.\",\n",
+    "        min=0.0,\n",
+    "        max=1.0,\n",
+    "    ),\n",
+    "    max_tokens: int = typer.Option(\n",
+    "        250, help=\"The maximum number of tokens to be used when generating a docstring for a function or class. Please note that a higher number will deplete your token quota faster.\"\n",
+    "    ),\n",
+    "    top_p: float = typer.Option(\n",
+    "        1.0, \n",
+    "        help=\"You can also specify a top-P value from 0-1 to achieve similar results to changing the temperature. According to the Open AI documentation, it is generally recommended to change either this or the temperature but not both.\",\n",
+    "        min=0.0,\n",
+    "        max=1.0,\n",
+    "   ),\n",
+    "    n: int = typer.Option(\n",
+    "        3, help=\"The number of docstrings to be generated for each function or class, with the best one being added to the source code. Please note that a higher number will deplete your token quota faster.\"\n",
+    "    ),\n",
+    ") -> None:\n",
+    "\n",
+    "    try:\n",
+    "        add_docstring_to_source(\n",
+    "            path=path,\n",
+    "            include_auto_gen_txt=include_auto_gen_txt,\n",
+    "            recreate_auto_gen_docs=recreate_auto_gen_docs,\n",
+    "            model=model,\n",
+    "            temperature=temperature,\n",
+    "            max_tokens=max_tokens,\n",
+    "            top_p=top_p,\n",
+    "            n=n,\n",
+    "        )\n",
+    "    except Exception as e:\n",
+    "        typer.secho(e, err=True, fg=typer.colors.RED)\n",
+    "        raise typer.Exit(1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9da7cdad",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">                                                                                                                   </span>\n",
+       "<span style=\"font-weight: bold\"> </span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">Usage: </span><span style=\"font-weight: bold\">gen [OPTIONS] [PATH]                                                                                       </span>\n",
+       "<span style=\"font-weight: bold\">                                                                                                                   </span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1m                                                                                                                   \u001b[0m\n",
+       "\u001b[1m \u001b[0m\u001b[1;33mUsage: \u001b[0m\u001b[1mgen [OPTIONS] [PATH]\u001b[0m\u001b[1m                                                                                      \u001b[0m\u001b[1m \u001b[0m\n",
+       "\u001b[1m                                                                                                                   \u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"> This command reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings  \n",
+       " to classes and methods that do not have them. It can be used to add missing docstrings to your codebase.          \n",
+       "                                                                                                                   \n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       " This command reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings  \n",
+       " to classes and methods that do not have them. It can be used to add missing docstrings to your codebase.          \n",
+       "                                                                                                                   \n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">╭─ Arguments ─────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>   path      <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f; font-weight: bold\">[</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">PATH</span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f; font-weight: bold\">]</span>  The path to the Jupyter notebook or Python file, or a directory containing these files      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: .]                                                                          </span>      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[2m╭─\u001b[0m\u001b[2m Arguments \u001b[0m\u001b[2m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[2m─╮\u001b[0m\n",
+       "\u001b[2m│\u001b[0m   path      \u001b[1;2;33m[\u001b[0m\u001b[1;33mPATH\u001b[0m\u001b[1;2;33m]\u001b[0m  The path to the Jupyter notebook or Python file, or a directory containing these files      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                     \u001b[2m[default: .]                                                                          \u001b[0m      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--prompt</span>                                                          <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">TEXT   </span>  Text that will be given as input to  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            the GPT-3 model to generate the      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            docstring. If no text is provided,   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            the docstring will be generated      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            according to the Google Python Style <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            Guide.                               <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: None]                     </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--include-auto-gen-txt</span>               <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">--no-include-auto-gen-txt</span>    <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">       </span>  If set to True, a note indicating    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            that the docstring was autogenerated <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            by docstring-gen library will be     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            added to the end.                    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: include-auto-gen-txt]     </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--force-recreate-auto-generated</span>  <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">-f</span>                               <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">       </span>  If set to True, the autogenerated    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            docstrings from the previous runs    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            will be replaced with the new one.   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--model</span>                                                           <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">TEXT   </span>  The name of the GPT-3 model to use   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            for docstring generation.            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: code-davinci-002]         </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--temperature</span>                                                     <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">FLOAT  </span>  The temperature parameter for the    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            GPT-3 model.                         <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 0.2]                      </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--max-tokens</span>                                                      <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">INTEGER</span>  The maximum number of tokens to      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            generate in the docstring.           <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 250]                      </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--top-p</span>                                                           <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">FLOAT  </span>  The top-p parameter for the GPT-3    <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            model.                               <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 1.0]                      </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--frequency-penalty</span>                                               <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">FLOAT  </span>  The frequency penalty parameter for  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            the GPT-3 model.                     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 0.0]                      </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--presence-penalty</span>                                                <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">FLOAT  </span>  The presence penalty parameter for   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            the GPT-3 model.                     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 0.0]                      </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--stop</span>                                                            <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">TEXT   </span>  A list of strings that, if           <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            encountered in the generated         <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            docstring, will cause generation to  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            stop.                                <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: #, \"\"\"]                   </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--n</span>                                                               <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">INTEGER</span>  The number of completions to         <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            generate for each prompt.            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">[default: 3]                        </span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--install-completion</span>                                              <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">       </span>  Install completion for the current   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            shell.                               <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--show-completion</span>                                                 <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">       </span>  Show completion for the current      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            shell, to copy it or customize the   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>                                                                            installation.                        <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">--help</span>                                                            <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">       </span>  Show this message and exit.          <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│</span>\n",
+       "<span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[2m╭─\u001b[0m\u001b[2m Options \u001b[0m\u001b[2m──────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\u001b[2m─╮\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-prompt\u001b[0m                                                          \u001b[1;33mTEXT   \u001b[0m  Text that will be given as input to  \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            the GPT-3 model to generate the      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            docstring. If no text is provided,   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            the docstring will be generated      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            according to the Google Python Style \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            Guide.                               \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: None]                     \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-include\u001b[0m\u001b[1;36m-auto-gen-txt\u001b[0m               \u001b[1;35m-\u001b[0m\u001b[1;35m-no\u001b[0m\u001b[1;35m-include-auto-gen-txt\u001b[0m    \u001b[1;33m       \u001b[0m  If set to True, a note indicating    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            that the docstring was autogenerated \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            by docstring-gen library will be     \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            added to the end.                    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: include-auto-gen-txt]     \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-force\u001b[0m\u001b[1;36m-recreate-auto-generated\u001b[0m  \u001b[1;32m-f\u001b[0m                               \u001b[1;33m       \u001b[0m  If set to True, the autogenerated    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            docstrings from the previous runs    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            will be replaced with the new one.   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-model\u001b[0m                                                           \u001b[1;33mTEXT   \u001b[0m  The name of the GPT-3 model to use   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            for docstring generation.            \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: code-davinci-002]         \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-temperature\u001b[0m                                                     \u001b[1;33mFLOAT  \u001b[0m  The temperature parameter for the    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            GPT-3 model.                         \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 0.2]                      \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-max\u001b[0m\u001b[1;36m-tokens\u001b[0m                                                      \u001b[1;33mINTEGER\u001b[0m  The maximum number of tokens to      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            generate in the docstring.           \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 250]                      \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-top\u001b[0m\u001b[1;36m-p\u001b[0m                                                           \u001b[1;33mFLOAT  \u001b[0m  The top-p parameter for the GPT-3    \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            model.                               \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 1.0]                      \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-frequency\u001b[0m\u001b[1;36m-penalty\u001b[0m                                               \u001b[1;33mFLOAT  \u001b[0m  The frequency penalty parameter for  \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            the GPT-3 model.                     \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 0.0]                      \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-presence\u001b[0m\u001b[1;36m-penalty\u001b[0m                                                \u001b[1;33mFLOAT  \u001b[0m  The presence penalty parameter for   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            the GPT-3 model.                     \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 0.0]                      \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-stop\u001b[0m                                                            \u001b[1;33mTEXT   \u001b[0m  A list of strings that, if           \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            encountered in the generated         \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            docstring, will cause generation to  \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            stop.                                \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: #, \"\"\"]                   \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-n\u001b[0m                                                               \u001b[1;33mINTEGER\u001b[0m  The number of completions to         \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            generate for each prompt.            \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            \u001b[2m[default: 3]                        \u001b[0m \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-install\u001b[0m\u001b[1;36m-completion\u001b[0m                                              \u001b[1;33m       \u001b[0m  Install completion for the current   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            shell.                               \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-show\u001b[0m\u001b[1;36m-completion\u001b[0m                                                 \u001b[1;33m       \u001b[0m  Show completion for the current      \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            shell, to copy it or customize the   \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m                                                                            installation.                        \u001b[2m│\u001b[0m\n",
+       "\u001b[2m│\u001b[0m \u001b[1;36m-\u001b[0m\u001b[1;36m-help\u001b[0m                                                            \u001b[1;33m       \u001b[0m  Show this message and exit.          \u001b[2m│\u001b[0m\n",
+       "\u001b[2m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "result = runner.invoke(_app, [\"--help\"])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0b4fe38",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/nbs/Docstring_Generator.ipynb b/nbs/Docstring_Generator.ipynb
index c1c6f9f..e596cf6 100644
--- a/nbs/Docstring_Generator.ipynb
+++ b/nbs/Docstring_Generator.ipynb
@@ -3,7 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2ae108fa",
+   "id": "24d0215b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8b32fc4f",
+   "id": "31594ead",
    "metadata": {},
    "source": [
     "# Docstring Generator"
@@ -21,573 +21,1697 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "63941987",
+   "id": "ac658e8e",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
+    "import time\n",
+    "import random\n",
     "import ast\n",
-    "import tokenize\n",
+    "import textwrap\n",
+    "import os\n",
+    "import re\n",
     "from typing import *\n",
     "from pathlib import Path\n",
-    "from io import BytesIO\n",
     "\n",
-    "import nbformat"
+    "import nbformat\n",
+    "import openai\n",
+    "import typer\n",
+    "\n",
+    "from mypy_extensions import NamedArg"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fab9b3ba",
+   "id": "0987791a",
    "metadata": {},
    "outputs": [],
    "source": [
     "import shutil\n",
-    "from tempfile import TemporaryDirectory"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "602a7fff",
-   "metadata": {},
-   "source": [
-    "## Docstring generator functions"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6cdabd36",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# | export\n",
+    "from tempfile import TemporaryDirectory\n",
+    "from contextlib import contextmanager\n",
+    "import unittest.mock\n",
     "\n",
-    "\n",
-    "def _generate_docstring_using_codex(code: str) -> str:\n",
-    "    return \"\"\"Sample docstring\n",
-    "\n",
-    "    Args:\n",
-    "        s: sample args\n",
-    "\n",
-    "    Returns:\n",
-    "        sample return\n",
-    "\"\"\""
+    "import pytest"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f0c01fc7",
+   "id": "1bce41e7",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
     "\n",
-    "def _inject_docstring_to_source(\n",
-    "    source: str, docstring: str, lineno: int, node_col_offset: int\n",
-    ") -> str:\n",
-    "    \"\"\"Inject a docstring into the source code at a specified line number.\n",
+    "def _visit_functions(\n",
+    "    tree: ast.AST,\n",
+    "    *,\n",
+    "    source: str,\n",
+    "    start_lineno: int = 1,\n",
+    "    end_lineno: int = -1,\n",
+    "    callback: Callable[\n",
+    "        [ast.AST, str, int, int, NamedArg(List[Tuple[int, int, int, int]], \"retval\")],\n",
+    "        Any,\n",
+    "    ],\n",
+    "    **kwargs: Any\n",
+    ") -> None:\n",
+    "    \"\"\"Walk the abstract syntax tree and call the callback function for every node found\n",
     "\n",
     "    Args:\n",
-    "        source: the source code\n",
-    "        docstring: the docstring to be added\n",
-    "        lineno: the line number at which the docstring will be inserted\n",
-    "        node_col_offset: the number of spaces to indent the docstring\n",
+    "        tree: The python AST\n",
+    "        source: The source code\n",
+    "        start_lineno: The start line number\n",
+    "        end_lineno: The end line number\n",
+    "        callback: The callback function\n",
+    "        kwargs: The keyword arguments\n",
     "\n",
-    "    Returns:\n",
-    "        The updated source code with the docstring injected\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
     "    \"\"\"\n",
-    "    lines = source.split(\"\\n\")\n",
-    "    indented_docstring = \"\\n\".join(\n",
-    "        [\n",
-    "            line\n",
-    "            if i == 0 or i == len(docstring.split(\"\\n\")) - 1\n",
-    "            else f\"{' ' * node_col_offset}{line}\"\n",
-    "            for i, line in enumerate(docstring.split(\"\\n\"))\n",
-    "        ]\n",
-    "    )\n",
-    "    indent = node_col_offset + 4\n",
-    "    lines.insert(lineno, f'{\" \" * indent}\"\"\"{indented_docstring}{\" \" * indent}\"\"\"')\n",
-    "    return \"\\n\".join(lines)"
+    "    if end_lineno == -1:\n",
+    "        end_lineno = len(source.split(\"\\n\"))\n",
+    "\n",
+    "    if isinstance(tree, (ast.AsyncFunctionDef, ast.FunctionDef, ast.ClassDef)):\n",
+    "        if ast.get_docstring(tree) is None:\n",
+    "            callback(tree, source, start_lineno, end_lineno, **kwargs)\n",
+    "    \n",
+    "    if hasattr(tree, \"body\"):\n",
+    "        for i, n in enumerate(tree.body):\n",
+    "            node_start_lineno = tree.body[i].lineno\n",
+    "            node_end_lineno = (\n",
+    "                tree.body[i + 1].lineno - 1 if i < (len(tree.body) - 1) else end_lineno\n",
+    "            )\n",
+    "            _visit_functions(\n",
+    "                n,\n",
+    "                source=source,\n",
+    "                start_lineno=node_start_lineno,\n",
+    "                end_lineno=node_end_lineno,\n",
+    "                callback=callback,\n",
+    "                **kwargs\n",
+    "            )"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3af581c1",
+   "id": "00a69169",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\n",
-      "async def drive(self):\n",
-      "    \"\"\"Sample docstring\n",
-      "\n",
-      "    Args:\n",
-      "        s: sample args\n",
-      "\n",
-      "    Returns:\n",
-      "        sample return\n",
-      "    \"\"\"\n",
-      "    print(f'The {self.model} is now driving.')\n",
-      "\n"
+      "[(2, 16), (5, 7), (8, 12), (17, 51)]\n"
      ]
     }
    ],
    "source": [
     "source = \"\"\"\n",
-    "async def drive(self):\n",
-    "    print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
+    "class Test:\n",
+    "    CONST_VAL = 1\n",
+    "    \n",
+    "    def __init__(self, a):\n",
+    "        self.a = a\n",
+    "        \n",
+    "    async def drive(self):\n",
+    "    \n",
+    "    \n",
+    "        print(f'The {self.model} is now driving.')\n",
+    "        \n",
+    "    def func_with_docstring():\n",
+    "        \\\"\"\" Sample docstring\\\"\"\"\n",
+    "        pass\n",
+    "        \n",
+    "def _check_and_add_docstrings_to_source(\n",
+    "    source: str, include_auto_gen_txt: bool, **kwargs\n",
+    ") -> str:\n",
+    "    source = _remove_auto_generated_docstring(source)    \n",
+    "    tree = ast.parse(source)\n",
+    "    line_offset = 0\n",
     "\n",
-    "docstring = \"\"\"Sample docstring\n",
+    "    for node in tree.body:\n",
+    "        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+    "            continue\n",
+    "        \n",
+    "        if ast.get_docstring(node) is not None:\n",
+    "            continue\n",
     "\n",
-    "    Args:\n",
-    "        s: sample args\n",
+    "        source, line_offset = _add_docstring(\n",
+    "            source, node, line_offset, include_auto_gen_txt, **kwargs\n",
+    "        )\n",
+    "        if not isinstance(node, ast.ClassDef):\n",
+    "            continue\n",
+    "        # Is a class and we need to check the functions inside\n",
+    "        # 29 - 36 make it as a recursive function\n",
+    "        for f in node.body:\n",
+    "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+    "                continue\n",
+    "            \n",
+    "            if ast.get_docstring(f) is not None:\n",
+    "                continue\n",
     "\n",
-    "    Returns:\n",
-    "        sample return\n",
-    "\"\"\"\n",
+    "            # should be a function inside the class for which there is no docstring\n",
+    "            source, line_offset = _add_docstring(\n",
+    "                source, f, line_offset, include_auto_gen_txt, **kwargs\n",
+    "            )\n",
     "\n",
-    "expected = \"\"\"\n",
-    "async def drive(self):\n",
-    "    \\\"\"\"Sample docstring\n",
+    "    return source\n",
+    "    \n",
+    "def test_callback(tree, source, start_lineno, end_lineno):\n",
+    "    \\\"\"\"Prints the source code of the given node\n",
     "\n",
     "    Args:\n",
-    "        s: sample args\n",
-    "\n",
-    "    Returns:\n",
-    "        sample return\n",
+    "        tree: The AST tree\n",
+    "        source: The source code\n",
+    "        start_lineno: The start line number\n",
+    "        end_lineno: The end line number\n",
     "    \\\"\"\"\n",
-    "    print(f'The {self.model} is now driving.')\n",
+    "\n",
+    "    pass\n",
     "\"\"\"\n",
+    "test_callback = unittest.mock.MagicMock()\n",
+    "tree = ast.parse(source)\n",
     "\n",
-    "lineno = 2\n",
-    "node_col_offset = 0\n",
-    "actual = _inject_docstring_to_source(source, docstring, lineno, node_col_offset)\n",
-    "print(actual)\n",
+    "_visit_functions(tree, source=source, callback=test_callback)\n",
     "\n",
-    "assert actual == expected"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a7b4e26e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# | export\n",
+    "actual = [(start_lineno, end_lineno) for tree, source, start_lineno, end_lineno in [x[1] for x in test_callback.mock_calls]]\n",
+    "expected = [(2, 16), (5, 7), (8, 12), (17, 51)]\n",
     "\n",
-    "def _get_code_from_source(source: str, start_line_no: int, end_line_no: int) -> str:\n",
-    "    \"\"\"Get a block of lines from a given source string.\n",
-    "    \n",
-    "    Args:\n",
-    "        source: The source string.\n",
-    "        start_line_no: The line number of the start of the block of lines.\n",
-    "        end_line_no: The line number of the end of the block of lines.\n",
-    "    \n",
-    "    Returns:\n",
-    "        The extracted block of lines from the source\n",
-    "    \"\"\"\n",
-    "    source_lines = source.split(\"\\n\")\n",
-    "    extracted_lines = source_lines[start_line_no-1:end_line_no]\n",
-    "    return \"\\n\".join(extracted_lines)"
+    "print(actual)\n",
+    "assert actual == expected, actual"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "710980d2",
+   "id": "eecd166b",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "    async def drive(self):\n",
-      "        print(f'The {self.model} is now driving.')\n"
+      "[]\n"
      ]
     }
    ],
    "source": [
-    "source = \"\"\"\n",
-    "class test:\n",
+    "source_with_dosctring = '''\n",
+    "class Test:\n",
+    "    \"\"\" Sample docstring\"\"\"\n",
     "    CONST_VAL = 1\n",
+    "    \n",
     "    def __init__(self, a):\n",
+    "        \"\"\" Sample docstring\"\"\"\n",
     "        self.a = a\n",
     "        \n",
     "    async def drive(self):\n",
+    "        \"\"\" Sample docstring\"\"\"\n",
     "        print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "start_line_no = 7\n",
-    "end_line_no = 8\n",
-    "\n",
-    "expected = \"\"\"    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\"\"\"\n",
-    "\n",
-    "actual = _get_code_from_source(source, start_line_no, end_line_no)\n",
-    "print(actual)\n",
+    "        \n",
+    "    def func_with_docstring():\n",
+    "        \"\"\" Sample docstring\"\"\"\n",
+    "        pass\n",
+    "'''\n",
+    "test_callback = unittest.mock.MagicMock()\n",
+    "tree = ast.parse(source_with_dosctring)\n",
     "\n",
-    "assert actual == expected\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d5526570",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# | export\n",
+    "_visit_functions(tree, source=source_with_dosctring, callback=test_callback)\n",
     "\n",
+    "actual = [(start_lineno, end_lineno) for tree, source, start_lineno, end_lineno in [x[1] for x in test_callback.mock_calls]]\n",
+    "expected = []\n",
     "\n",
-    "def _calculate_end_lineno(source: str, start_line_no: int) -> int:\n",
-    "    \"\"\"Calculate the end line number of a function in a Python source code.\n",
-    "    \n",
-    "    Args:\n",
-    "        source: The source code string.\n",
-    "        start_line_no: The line number of the start of the function.\n",
-    "        \n",
-    "    Returns:\n",
-    "        The end line number of the function.\n",
-    "    \"\"\"\n",
-    "    lines = source.split(\"\\n\")[start_line_no - 1 :]\n",
-    "    first_indent = len(lines[0]) - len(lines[0].lstrip())\n",
-    "    end_line_in_source = 0\n",
-    "\n",
-    "    for i, line in enumerate(lines[1:]):\n",
-    "        if len(line) - len(line.lstrip()) == first_indent and line.strip() != \"\":\n",
-    "            end_line_in_source = i\n",
-    "            break\n",
-    "\n",
-    "    ret_val = (\n",
-    "        len(source.split(\"\\n\"))\n",
-    "        if end_line_in_source == 0\n",
-    "        else end_line_in_source + start_line_no\n",
-    "    )\n",
-    "    return ret_val - 1"
+    "print(actual)\n",
+    "assert actual == expected, actual"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "18c65bc0",
+   "id": "9494c957",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[8, 5, 9, 19, 13, 16, 20, 27, 24, 28, 30]\n"
+      "************************************************************************************************************************\n",
+      "class Test:\n",
+      "    CONST_VAL = 1\n",
+      "    \n",
+      "    def __init__(self, a):\n",
+      "        self.a = a\n",
+      "        \n",
+      "    async def drive(self):\n",
+      "    \n",
+      "    \n",
+      "        print(f'The {self.model} is now driving.')\n",
+      "        \n",
+      "    def func_with_docstring():\n",
+      "        \"\"\" Sample docstring\"\"\"\n",
+      "        pass\n",
+      "        \n",
+      "************************************************************************************************************************\n",
+      "    def __init__(self, a):\n",
+      "        self.a = a\n",
+      "        \n",
+      "************************************************************************************************************************\n",
+      "    async def drive(self):\n",
+      "    \n",
+      "    \n",
+      "        print(f'The {self.model} is now driving.')\n",
+      "        \n",
+      "************************************************************************************************************************\n",
+      "def _check_and_add_docstrings_to_source(\n",
+      "    source: str, include_auto_gen_txt: bool, **kwargs\n",
+      ") -> str:\n",
+      "    source = _remove_auto_generated_docstring(source)    \n",
+      "    tree = ast.parse(source)\n",
+      "    line_offset = 0\n",
+      "\n",
+      "    for node in tree.body:\n",
+      "        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "            continue\n",
+      "        \n",
+      "        if ast.get_docstring(node) is not None:\n",
+      "            continue\n",
+      "\n",
+      "        source, line_offset = _add_docstring(\n",
+      "            source, node, line_offset, include_auto_gen_txt, **kwargs\n",
+      "        )\n",
+      "        if not isinstance(node, ast.ClassDef):\n",
+      "            continue\n",
+      "        # Is a class and we need to check the functions inside\n",
+      "        # 29 - 36 make it as a recursive function\n",
+      "        for f in node.body:\n",
+      "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "                continue\n",
+      "            \n",
+      "            if ast.get_docstring(f) is not None:\n",
+      "                continue\n",
+      "\n",
+      "            # should be a function inside the class for which there is no docstring\n",
+      "            source, line_offset = _add_docstring(\n",
+      "                source, f, line_offset, include_auto_gen_txt, **kwargs\n",
+      "            )\n",
+      "\n",
+      "    return source\n",
+      "    \n"
      ]
     }
    ],
    "source": [
-    "source = \"\"\"\n",
-    "class A:\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        self.a = a\n",
-    "        \n",
-    "    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\n",
     "\n",
-    "class B:\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        self.a = a\n",
-    "        \n",
-    "    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\n",
-    "        \n",
-    "    async def stop(self):\n",
-    "        print(f'The {self.model} is now stopped.')\n",
-    "        \n",
-    "class C:\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        self.a = a\n",
+    "def test_callback(tree, source, start_lineno, end_lineno):\n",
+    "    \"\"\"Prints the source code of the given node\n",
     "\n",
-    "    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\n",
-    "        \n",
-    "async def drive(self):\n",
-    "    print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
+    "    Args:\n",
+    "        tree: The AST tree\n",
+    "        source: The source code\n",
+    "        start_lineno: The start line number\n",
+    "        end_lineno: The end line number\n",
+    "\n",
+    "    \"\"\"\n",
+    "\n",
+    "    source_lines = source.split(\"\\n\")\n",
+    "    node_lines = source_lines[start_lineno - 1 : end_lineno]\n",
+    "    print(\"*\" * 120)\n",
+    "    print(\"\\n\".join(node_lines))\n",
     "\n",
+    "    \n",
     "tree = ast.parse(source)\n",
-    "expected = [8, 5, 9, 19, 13, 16, 20, 27, 24, 28, 30]\n",
-    "actual = []\n",
-    "for i, node in enumerate(tree.body):\n",
-    "    res = _calculate_end_lineno(source, node.lineno)\n",
-    "    actual.append(res)\n",
-    "#     print(node.end_lineno)\n",
-    "\n",
-    "    if isinstance(node, ast.ClassDef):\n",
-    "        for f in node.body:\n",
-    "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
-    "                continue\n",
-    "            res = _calculate_end_lineno(source, f.lineno)\n",
-    "            actual.append(res)\n",
-    "#             print(f.end_lineno)\n",
-    "print(actual)\n",
-    "assert actual == expected"
+    "\n",
+    "_visit_functions(tree, source=source, callback=test_callback)\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "acae29db",
+   "id": "87aff27c",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
     "\n",
-    "def _line_has_class_or_method(source: str, lineno: int) -> bool:\n",
-    "    \"\"\"Check if a line in the source code contains a class or method definition.\n",
-    "    \n",
-    "    Args:\n",
-    "        source: The source code as a string.\n",
-    "        lineno: The line number to check.\n",
-    "    \n",
-    "    Returns:\n",
-    "        True if the line contains a class or method definition, False otherwise.\n",
-    "    \"\"\"\n",
-    "    line = \"\".join(source.split(\"\\n\")[lineno - 1])\n",
-    "    tokens = list(tokenize.tokenize(BytesIO(line.encode(\"utf-8\")).readline))\n",
-    "    return tokens[1].type == tokenize.NAME and tokens[1].string in {\n",
-    "        \"class\",\n",
-    "        \"def\",\n",
-    "        \"async\",\n",
-    "    }\n",
-    "\n",
-    "def _get_start_line_for_class_or_func(source: str, lineno: int) -> int:\n",
-    "    \"\"\"Get the line number of the first line containing a class or function definition.\n",
-    "    \n",
+    "def _get_classes_and_functions(source: str) -> List[Tuple[int, int, int, int]]:\n",
+    "    \"\"\"Get the classes and functions in a source file.\n",
+    "\n",
     "    Args:\n",
-    "        source: The source code as a string.\n",
-    "        lineno: The line number to start from.\n",
-    "    \n",
+    "        source: The source code of the file.\n",
+    "        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\n",
+    "\n",
+    "\n",
     "    Returns:\n",
-    "        The line number of the first line containing a class or function definition, \n",
-    "        or the original line number if no such line is found.\n",
+    "        A list of tuples of the form (start_lineno, end_lineno, start_col_offset, end_col_offset)\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
     "    \"\"\"\n",
-    "    if _line_has_class_or_method(source, lineno):\n",
-    "        return lineno\n",
-    "\n",
-    "    original_lineno = lineno\n",
-    "    total_lines = source.split(\"\\n\")\n",
-    "    for i in total_lines:\n",
-    "        lineno += 1\n",
-    "        if lineno > len(total_lines):\n",
-    "            break\n",
-    "        if _line_has_class_or_method(source, lineno):\n",
-    "            return lineno\n",
-    "    return original_lineno"
+    "    retval: List[Tuple[int, int, int, int]] = []\n",
+    "    tree = ast.parse(source)\n",
+    "\n",
+    "    def callback(tree, source, start_lineno, end_lineno, *, retval):\n",
+    "        \"\"\"Callback function for the ast.walk function.\n",
+    "\n",
+    "        Args:\n",
+    "            tree: The tree to walk\n",
+    "            source: The source code\n",
+    "            start_lineno: The starting line number\n",
+    "            end_lineno: The ending line number\n",
+    "\n",
+    "        Keyword Args:\n",
+    "            retval: The return value\n",
+    "\n",
+    "        Returns:\n",
+    "            The return value\n",
+    "\n",
+    "        !!! note\n",
+    "\n",
+    "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "        \"\"\"\n",
+    "        retval.append(\n",
+    "            (start_lineno, tree.body[0].lineno - 1, end_lineno, tree.body[0].col_offset)\n",
+    "        )\n",
+    "\n",
+    "    _visit_functions(\n",
+    "        tree,\n",
+    "        source=source,\n",
+    "        callback=callback,\n",
+    "        retval=retval,\n",
+    "    )\n",
+    "\n",
+    "    return retval"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e0d08807",
+   "id": "72e81d66",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "30\n"
+      "[(2, 2, 16, 4), (5, 5, 7, 8), (8, 10, 12, 8), (17, 19, 51, 4)]\n"
      ]
     }
    ],
    "source": [
-    "source = \"\"\"def decorator1(func):\n",
-    "    \\\"\"\"Sample docstring\n",
-    "\n",
-    "    Args:\n",
-    "        s: sample args\n",
+    "actual = _get_classes_and_functions(source)\n",
+    "expected = [(2, 2, 16, 4), (5, 5, 7, 8), (8, 10, 12, 8), (17, 19, 51, 4)]\n",
+    "print(actual)\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "306d1b6e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
     "\n",
-    "    Returns:\n",
-    "        sample return\n",
-    "    \\\"\"\"\n",
-    "    def inner():\n",
-    "        func()\n",
-    "    return inner\n",
     "\n",
-    "def decorator2(func):\n",
-    "    \\\"\"\"Sample docstring\n",
+    "def _get_code_from_source(source: str, start_line_no: int, end_line_no: int) -> str:\n",
+    "    \"\"\"Get code from source\n",
     "\n",
     "    Args:\n",
-    "        s: sample args\n",
+    "        source: The source code of the file.\n",
+    "        start_line_no: Start line number\n",
+    "        end_line_no: End line number\n",
     "\n",
     "    Returns:\n",
-    "        sample return\n",
-    "    \\\"\"\"\n",
-    "    def inner():\n",
-    "        func()\n",
-    "    return inner\n",
+    "        The extracted code\n",
     "\n",
+    "    !!! note\n",
     "\n",
-    "@decorator1\n",
-    "@decorator2\n",
-    "def outer_func():\n",
-    "    def inner_func():\n",
-    "        print(\"Hello, World!\")\n",
-    "    inner_func()\n",
-    "\"\"\"\n",
-    "\n",
-    "line_no = 28\n",
-    "expected = 30\n",
-    "actual = _get_start_line_for_class_or_func(source, line_no)\n",
-    "print(actual)\n",
-    "assert actual == expected"
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    source_lines = source.split(\"\\n\")\n",
+    "    extracted_lines = source_lines[start_line_no - 1 : end_line_no]\n",
+    "    return \"\\n\".join(extracted_lines)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0e963690",
+   "id": "9c50770e",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "2\n",
-      "4\n",
-      "3\n",
-      "5\n",
-      "2\n"
-     ]
-    }
-   ],
-   "source": [
-    "source = \"\"\"\n",
-    "async def drive(self):\n",
-    "    print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "lineno = 2\n",
-    "expected = 2\n",
-    "actual = _get_start_line_for_class_or_func(source, lineno)\n",
-    "print(actual)\n",
-    "\n",
-    "assert actual == expected\n",
-    "\n",
-    "source = \"\"\"\n",
-    "@foo(x=5)\n",
-    "@decorator2\n",
-    "async def drive(self):\n",
-    "    print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "lineno = 2\n",
-    "expected = 4\n",
-    "actual = _get_start_line_for_class_or_func(source, lineno)\n",
-    "print(actual)\n",
-    "\n",
-    "assert actual == expected\n",
+      "[(2, 2, 16, 4), (5, 5, 7, 8), (8, 10, 12, 8), (17, 19, 51, 4)]\n",
+      "****************************************************************************************************\n",
+      "class Test:\n",
+      "    CONST_VAL = 1\n",
+      "    \n",
+      "    def __init__(self, a):\n",
+      "        self.a = a\n",
+      "        \n",
+      "    async def drive(self):\n",
+      "    \n",
+      "    \n",
+      "        print(f'The {self.model} is now driving.')\n",
+      "        \n",
+      "    def func_with_docstring():\n",
+      "        \"\"\" Sample docstring\"\"\"\n",
+      "        pass\n",
+      "        \n",
+      "****************************************************************************************************\n",
+      "    def __init__(self, a):\n",
+      "        self.a = a\n",
+      "        \n",
+      "****************************************************************************************************\n",
+      "    async def drive(self):\n",
+      "    \n",
+      "    \n",
+      "        print(f'The {self.model} is now driving.')\n",
+      "        \n",
+      "****************************************************************************************************\n",
+      "def _check_and_add_docstrings_to_source(\n",
+      "    source: str, include_auto_gen_txt: bool, **kwargs\n",
+      ") -> str:\n",
+      "    source = _remove_auto_generated_docstring(source)    \n",
+      "    tree = ast.parse(source)\n",
+      "    line_offset = 0\n",
+      "\n",
+      "    for node in tree.body:\n",
+      "        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "            continue\n",
+      "        \n",
+      "        if ast.get_docstring(node) is not None:\n",
+      "            continue\n",
+      "\n",
+      "        source, line_offset = _add_docstring(\n",
+      "            source, node, line_offset, include_auto_gen_txt, **kwargs\n",
+      "        )\n",
+      "        if not isinstance(node, ast.ClassDef):\n",
+      "            continue\n",
+      "        # Is a class and we need to check the functions inside\n",
+      "        # 29 - 36 make it as a recursive function\n",
+      "        for f in node.body:\n",
+      "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "                continue\n",
+      "            \n",
+      "            if ast.get_docstring(f) is not None:\n",
+      "                continue\n",
+      "\n",
+      "            # should be a function inside the class for which there is no docstring\n",
+      "            source, line_offset = _add_docstring(\n",
+      "                source, f, line_offset, include_auto_gen_txt, **kwargs\n",
+      "            )\n",
+      "\n",
+      "    return source\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "linenos = _get_classes_and_functions(source)\n",
+    "print(linenos)\n",
+    "actual = [\n",
+    "    _get_code_from_source(source, start_line_no, end_line_no)\n",
+    "    for start_line_no, docstring_line_no, end_line_no, node_offset in linenos\n",
+    "]\n",
+    "expected = [\n",
+    "    'class Test:\\n    CONST_VAL = 1\\n    \\n    def __init__(self, a):\\n        self.a = a\\n        \\n    async def drive(self):\\n    \\n    \\n        print(f\\'The {self.model} is now driving.\\')\\n        \\n    def func_with_docstring():\\n        \"\"\" Sample docstring\"\"\"\\n        pass\\n        ',\n",
+    "    \"    def __init__(self, a):\\n        self.a = a\\n        \",\n",
+    "    \"    async def drive(self):\\n    \\n    \\n        print(f'The {self.model} is now driving.')\\n        \",\n",
+    "    \"def _check_and_add_docstrings_to_source(\\n    source: str, include_auto_gen_txt: bool, **kwargs\\n) -> str:\\n    source = _remove_auto_generated_docstring(source)    \\n    tree = ast.parse(source)\\n    line_offset = 0\\n\\n    for node in tree.body:\\n        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\\n            continue\\n        \\n        if ast.get_docstring(node) is not None:\\n            continue\\n\\n        source, line_offset = _add_docstring(\\n            source, node, line_offset, include_auto_gen_txt, **kwargs\\n        )\\n        if not isinstance(node, ast.ClassDef):\\n            continue\\n        # Is a class and we need to check the functions inside\\n        # 29 - 36 make it as a recursive function\\n        for f in node.body:\\n            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\\n                continue\\n            \\n            if ast.get_docstring(f) is not None:\\n                continue\\n\\n            # should be a function inside the class for which there is no docstring\\n            source, line_offset = _add_docstring(\\n                source, f, line_offset, include_auto_gen_txt, **kwargs\\n            )\\n\\n    return source\\n    \",\n",
+    "]\n",
+    "for f in actual:\n",
+    "    print(\"*\" * 100)\n",
+    "    print(f)\n",
+    "\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6ec43971",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "# Reference: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_handle_rate_limits.ipynb\n",
+    "\n",
+    "\n",
+    "def _retry_with_exponential_backoff(\n",
+    "    initial_delay: float = 1,\n",
+    "    exponential_base: float = 2,\n",
+    "    jitter: bool = True,\n",
+    "    max_retries: int = 10,\n",
+    "    max_wait: float = 60,\n",
+    "    errors: tuple = (\n",
+    "        openai.error.RateLimitError,\n",
+    "        openai.error.ServiceUnavailableError,\n",
+    "        openai.error.APIError,\n",
+    "    ),\n",
+    ") -> Callable:\n",
+    "    \"\"\"Retry a function with exponential backoff.\"\"\"\n",
+    "\n",
+    "    def decorator(func):\n",
+    "        \"\"\"Decorator to retry a function if it fails.\n",
+    "\n",
+    "        Args:\n",
+    "            func: The function to be decorated\n",
+    "            max_retries: The maximum number of retries\n",
+    "            initial_delay: The initial delay\n",
+    "            exponential_base: The exponential base\n",
+    "            max_wait: The maximum wait\n",
+    "            jitter: The jitter\n",
+    "\n",
+    "        Returns:\n",
+    "            The decorated function\n",
+    "\n",
+    "        Raises:\n",
+    "            Exception: If the maximum number of retries is exceeded\n",
+    "\n",
+    "        !!! note\n",
+    "\n",
+    "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "        \"\"\"\n",
+    "        def wrapper(*args, **kwargs):\n",
+    "            num_retries = 0\n",
+    "            delay = initial_delay\n",
+    "\n",
+    "            while True:\n",
+    "                try:\n",
+    "                    return func(*args, **kwargs)\n",
+    "\n",
+    "                except errors as e:\n",
+    "                    num_retries += 1\n",
+    "                    if num_retries > max_retries:\n",
+    "                        raise Exception(\n",
+    "                            f\"Maximum number of retries ({max_retries}) exceeded.\"\n",
+    "                        )\n",
+    "                    delay = min(\n",
+    "                        delay\n",
+    "                        * exponential_base\n",
+    "                        * (1 + jitter * random.random()),  # nosec\n",
+    "                        max_wait,\n",
+    "                    )\n",
+    "                    typer.secho(\n",
+    "                        f\"Note: OpenAI's API rate limit reached. Command will automatically retry in {int(delay)} seconds. For more information visit: https://help.openai.com/en/articles/5955598-is-api-usage-subject-to-any-rate-limits\",\n",
+    "                        fg=typer.colors.BLUE,\n",
+    "                    )\n",
+    "                    time.sleep(delay)\n",
+    "\n",
+    "                except Exception as e:\n",
+    "                    raise e\n",
+    "\n",
+    "        return wrapper\n",
+    "\n",
+    "    return decorator\n",
+    "\n",
+    "\n",
+    "@_retry_with_exponential_backoff()\n",
+    "def _completions_with_backoff(*args, **kwargs):\n",
+    "    return openai.Completion.create(*args, **kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51f5d9ba",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "DEFAULT_PROMPT_TEMPLATE = '''\n",
+    "# Python 3.7\n",
+    "\n",
+    "#docstring_gen: Original code:\n",
+    "\n",
+    "def add_strings(s1: Optional[str] = None, s2: Optional[str] = None) -> str:\n",
+    "    if s1 is None or s2 is None:\n",
+    "        raise ValueError(Both s1 and s2 must be provided and must be of type string\")\n",
+    "    return s1 + s2\n",
+    "    \n",
+    "#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.\n",
+    "\n",
+    "\"\"\"Add two strings\n",
+    "\n",
+    "Args:\n",
+    "    s1: First string\n",
+    "    s2: Second string\n",
+    "    \n",
+    "Returns:\n",
+    "    The added string\n",
+    "    \n",
+    "Raises:\n",
+    "    ValueError: If s1 or s2 is None\n",
+    "\"\"\"\n",
+    "\n",
+    "#docstring_gen: Original code:\n",
+    "\n",
+    "class Person:\n",
+    "    def __init__(self, name, surname, age):\n",
+    "        self.name = name\n",
+    "        self.surname = surname\n",
+    "        self.age = age\n",
+    "\n",
+    "    def info(self, additional=\"\"):\n",
+    "        print('My name is :' + self.name + additional)\n",
+    "        \n",
+    "#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.\n",
+    "\n",
+    "\n",
+    "\"\"\"A class to represent a person.\n",
+    "\n",
+    "Attributes:\n",
+    "    name : first name of the person\n",
+    "    surname : family name of the person\n",
+    "    age : age of the person\n",
+    "\"\"\"\n",
+    "\n",
+    "#docstring_gen: Original code:\n",
+    "\n",
+    "{code}\n",
+    "\n",
+    "#docstring_gen: A comprehensive PEP 257 Google style doctring, including a brief one-line summary of the function.\n",
+    "\n",
+    "\"\"\"\n",
+    "'''"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a18a7daa",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{\n",
+      "  \"choices\": [\n",
+      "    {\n",
+      "      \"finish_reason\": \"stop\",\n",
+      "      \"index\": 0,\n",
+      "      \"logprobs\": null,\n",
+      "      \"text\": \"Drive a car.\\n\\nArgs:\\n    name: The name of the car.\\n\\nReturns:\\n    None\\n\"\n",
+      "    }\n",
+      "  ],\n",
+      "  \"created\": 1674027647,\n",
+      "  \"id\": \"cmpl-6ZxLrWeHCILfnlvwA81jIIAaDJkB4\",\n",
+      "  \"model\": \"code-davinci-002\",\n",
+      "  \"object\": \"text_completion\",\n",
+      "  \"usage\": {\n",
+      "    \"completion_tokens\": 27,\n",
+      "    \"prompt_tokens\": 397,\n",
+      "    \"total_tokens\": 424\n",
+      "  }\n",
+      "}\n"
+     ]
+    }
+   ],
+   "source": [
+    "code = \"\"\"\n",
+    "def drive(name: str) -> None:\n",
+    "    print(f'The {name} is now driving.')\n",
+    "\"\"\"\n",
+    "\n",
+    "prompt = DEFAULT_PROMPT_TEMPLATE.format(code=code)\n",
+    "\n",
+    "response = _completions_with_backoff(\n",
+    "    prompt=prompt,\n",
+    "    model=\"code-davinci-002\",\n",
+    "    temperature=0.2,\n",
+    "    max_tokens=250,\n",
+    "    top_p=1.0,\n",
+    "    frequency_penalty=0.0,\n",
+    "    presence_penalty=0.0,\n",
+    "    stop=['\"\"\"'],\n",
+    "    n=1,\n",
+    ")\n",
+    "\n",
+    "print(response)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5ebb240",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "@contextmanager\n",
+    "def mock_openai_create():\n",
+    "    mock_choice = unittest.mock.MagicMock\n",
+    "    mock_choice.text = \"\"\"Drive a car\n",
+    "\n",
+    "    Args:\n",
+    "        name: The name of the car\n",
+    "    \"\"\"\n",
+    "    mock_choices = unittest.mock.MagicMock\n",
+    "    mock_choices.choices = [mock_choice]\n",
+    "\n",
+    "    with unittest.mock.patch(\"openai.Completion\") as MockClass:\n",
+    "        MockClass.create.return_value = mock_choices\n",
+    "        yield\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6aa120e5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "with mock_openai_create():\n",
+    "    response = openai.Completion.create()\n",
+    "    print(response.choices[0].text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6ab98687",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mNote: OpenAI's API rate limit reached. Command will automatically retry in 2 seconds. For more information visit: https://help.openai.com/en/articles/5955598-is-api-usage-subject-to-any-rate-limits\u001b[0m\n",
+      "Maximum number of retries (1) exceeded.\n"
+     ]
+    }
+   ],
+   "source": [
+    "@_retry_with_exponential_backoff()\n",
+    "def mock_func():\n",
+    "    return \"Success\"\n",
+    "\n",
+    "assert mock_func() == \"Success\"\n",
+    "\n",
+    "# Test max retries exceeded\n",
+    "@_retry_with_exponential_backoff(max_retries=1)\n",
+    "def mock_func_error():\n",
+    "    raise openai.error.RateLimitError\n",
+    "\n",
+    "\n",
+    "with pytest.raises(Exception) as e:\n",
+    "    mock_func_error()\n",
+    "\n",
+    "print(e.value)\n",
+    "assert str(e.value) == \"Maximum number of retries (1) exceeded.\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00599e6c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "\n",
+    "def _get_best_docstring(docstrings: List[str]) -> Optional[str]:\n",
+    "    \"\"\"Get the best docstring from a list of docstrings.\n",
+    "\n",
+    "    Args:\n",
+    "        docstrings: List of docstrings\n",
+    "\n",
+    "    Returns:\n",
+    "        The best docstring\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    docstrings = [d for d in docstrings if \"Args:\" in d or \"Attributes:\" in d]\n",
+    "    docstrings = [d for d in docstrings if \"~~~~\" not in d]\n",
+    "    return docstrings[0] if len(docstrings) > 0 else None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8233ac5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "    This function checks if the source code has docstrings for all the functions and classes.\n",
+      "    If not, it adds the docstrings.\n",
+      "    It also removes the auto generated docstring.\n",
+      "    \n",
+      "    Args:\n",
+      "        source: The source code as a string.\n",
+      "        include_auto_gen_txt: Whether to include the auto generated text in the docstring.\n",
+      "        **kwargs: Other keyword arguments.\n",
+      "    \n",
+      "    Returns:\n",
+      "        The source code with docstrings added.\n",
+      "    \n",
+      "    Raises:\n",
+      "        ValueError: If the source code is not a string.\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "docstrings = [\n",
+    "    \"    _check_and_add_docstrings_to_source(\\n    source: str, include_auto_gen_txt: bool, **kwargs\\n) -> str:\\n    source = _remove_auto_generated_docstring(source)    \\n    tree = ast.parse(source)\\n    line_offset = 0\\n\\n    for node in tree.body:\\n        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\\n            continue\\n        \\n        if ast.get_docstring(node) is not None:\\n            continue\\n\\n        source, line_offset = _add_docstring(\\n            source, node, line_offset, include_auto_gen_txt, **kwargs\\n        )\\n        if not isinstance(node, ast.ClassDef):\\n            continue\\n        \",\n",
+    "    '    _check_and_add_docstrings_to_source\\n    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n    This function checks if the source code has docstrings for all the functions and classes.\\n    If not, it adds a docstring to the function/class.\\n\\n    Args:\\n        source (str): The source code to be checked for docstrings.\\n        include_auto_gen_txt (bool): If True, the docstring will include the text \"Auto-generated by nbdev\".\\n        **kwargs: Additional keyword arguments.\\n\\n    Returns:\\n        str: The source code with docstrings added.\\n',\n",
+    "    \"    This function checks if the source code has docstrings for all the functions and classes.\\n    If not, it adds the docstrings.\\n    It also removes the auto generated docstring.\\n    \\n    Args:\\n        source: The source code as a string.\\n        include_auto_gen_txt: Whether to include the auto generated text in the docstring.\\n        **kwargs: Other keyword arguments.\\n    \\n    Returns:\\n        The source code with docstrings added.\\n    \\n    Raises:\\n        ValueError: If the source code is not a string.\\n\",\n",
+    "]\n",
+    "\n",
+    "actual = _get_best_docstring(docstrings)\n",
+    "expected = docstrings[2]\n",
+    "\n",
+    "print(actual)\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "600e9458",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "None\n"
+     ]
+    }
+   ],
+   "source": [
+    "docstrings = [\n",
+    "    \"    _check_and_add_docstrings_to_source(\\n    source: str, include_auto_gen_txt: bool, **kwargs\\n) -> str:\\n    source = _remove_auto_generated_docstring(source)    \\n    tree = ast.parse(source)\\n    line_offset = 0\\n\\n    for node in tree.body:\\n        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\\n            continue\\n        \\n        if ast.get_docstring(node) is not None:\\n            continue\\n\\n        source, line_offset = _add_docstring(\\n            source, node, line_offset, include_auto_gen_txt, **kwargs\\n        )\\n        if not isinstance(node, ast.ClassDef):\\n            continue\\n        \",\n",
+    "    '    _check_and_add_docstrings_to_source\\n    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n    This function checks if the source code has docstrings for all the functions and classes.\\n    If not, it adds a docstring to the function/class.\\n\\n    Args:\\n        source (str): The source code to be checked for docstrings.\\n        include_auto_gen_txt (bool): If True, the docstring will include the text \"Auto-generated by nbdev\".\\n        **kwargs: Additional keyword arguments.\\n\\n    Returns:\\n        str: The source code with docstrings added.\\n',\n",
+    "    \"    This function checks if the source code ~~~~~~~~~~~~~~~~ has docstrings for all the functions and classes.\\n    If not, it adds the docstrings.\\n    It also removes the auto generated docstring.\\n    \\n    Args:\\n        source: The source code as a string.\\n        include_auto_gen_txt: Whether to include the auto generated text in the docstring.\\n        **kwargs: Other keyword arguments.\\n    \\n    Returns:\\n        The source code with docstrings added.\\n    \\n    Raises:\\n        ValueError: If the source code is not a string.\\n\",\n",
+    "]\n",
+    "\n",
+    "actual = _get_best_docstring(docstrings)\n",
+    "expected = None\n",
+    "\n",
+    "print(actual)\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23bc603f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "DOCSTRING_ERR = \"\"\"!!! note\n",
+    "    \n",
+    "    Failed to generate docs\n",
+    "\"\"\"\n",
+    "\n",
+    "\n",
+    "def _generate_docstring_using_codex(\n",
+    "    code: str,\n",
+    "    prompt_template: str,\n",
+    "    **kwargs: Union[int, float, str, List[str]]\n",
+    ") -> str:\n",
+    "    \"\"\"Generate a docstring using codex.\n",
+    "\n",
+    "    Args:\n",
+    "        code: The code for which to generate a docstring.\n",
+    "        prompt_template: The template for the prompt.\n",
+    "        **kwargs: The keyword arguments.\n",
+    "\n",
+    "    Returns:\n",
+    "        The generated docstring.\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    prompt = prompt_template.format(code=code)\n",
+    "    response = _completions_with_backoff(prompt=prompt, **kwargs)\n",
+    "    \n",
+    "    best_docstring = _get_best_docstring([i.text for i in response.choices])\n",
+    "    retval = best_docstring if best_docstring is not None else DOCSTRING_ERR\n",
+    "    return retval"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76f27b70",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "    \n",
+      "Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "    \n",
+      "Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "    \n",
+      "Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "    \n"
+     ]
+    }
+   ],
+   "source": [
+    "linenos = _get_classes_and_functions(source)\n",
+    "classes_and_functions = [\n",
+    "    _get_code_from_source(source, start_line_no, end_line_no)\n",
+    "    for start_line_no, docstring_line_no, end_line_no, node_offset in linenos\n",
+    "]\n",
+    "with mock_openai_create():\n",
+    "    for code in classes_and_functions:\n",
+    "        docstring = _generate_docstring_using_codex(\n",
+    "            code,\n",
+    "            DEFAULT_PROMPT_TEMPLATE,\n",
+    "            model=\"code-davinci-002\",\n",
+    "            temperature=0.2,\n",
+    "            max_tokens=250,\n",
+    "            top_p=1.0,\n",
+    "            frequency_penalty=0.0,\n",
+    "            presence_penalty=0.0,\n",
+    "            stop=['\"\"\"'],\n",
+    "            n=3,\n",
+    "        )\n",
+    "\n",
+    "        print(docstring)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "797503b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "AUTO_GEN_PERFIX = \"\"\"\n",
+    "!!! note\n",
+    "\n",
+    "\"\"\"\n",
+    "AUTO_GEN_BODY = \"The above docstring is autogenerated by docstring-gen library\"\n",
+    "AUTO_GEN_SUFFIX = \"(https://github.com/airtai/docstring-gen)\"\n",
+    "\n",
+    "AUTO_GEN_TXT = AUTO_GEN_PERFIX + \" \" * 4 + AUTO_GEN_BODY + \" \" + AUTO_GEN_SUFFIX\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f43a8bb6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "def _add_auto_gen_txt(docstring: str) -> str:\n",
+    "    \"\"\"Add the autogenerated by docstring-gen library text to the end of a docstring\n",
+    "\n",
+    "    Args:\n",
+    "        docstring: The docstring to which the text will be added\n",
+    "\n",
+    "    Returns:\n",
+    "        The docstring with the added text\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    return docstring + AUTO_GEN_TXT + \"\\n\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59e7ce75",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sample docstring\n",
+      "\n",
+      "Args:\n",
+      "    s: sample args\n",
+      "\n",
+      "Returns:\n",
+      "    sample return\n",
+      "\n",
+      "!!! note\n",
+      "\n",
+      "    The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "docstring = '''Sample docstring\n",
+    "\n",
+    "Args:\n",
+    "    s: sample args\n",
+    "\n",
+    "Returns:\n",
+    "    sample return\n",
+    "'''\n",
+    "\n",
+    "expected = docstring + AUTO_GEN_TXT + \"\\n\"\n",
+    "\n",
+    "actual = _add_auto_gen_txt(docstring)\n",
+    "print(actual)\n",
+    "\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39e67bf6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "\n",
+    "def _fix_docstring_indent(\n",
+    "    docstring: str, col_offset: int, *, include_auto_gen_txt: bool\n",
+    ") -> str:\n",
+    "    \"\"\"Fix the indentation of a docstring.\n",
+    "\n",
+    "    Args:\n",
+    "        docstring: The docstring to fix.\n",
+    "        col_offset: The column offset to use.\n",
+    "        include_auto_gen_txt: If True, include the auto-generated text in the docstring.\n",
+    "\n",
+    "    Returns:\n",
+    "        The fixed docstring.\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    lines = docstring.split(\"\\n\")\n",
+    "    rest = textwrap.dedent(\"\\n\".join(lines[1:]))\n",
+    "    retval = lines[0] + \"\\n\" + rest\n",
+    "\n",
+    "    if include_auto_gen_txt:\n",
+    "        retval = _add_auto_gen_txt(retval)\n",
+    "    \n",
+    "    retval = '\"\"\"' + retval + '\"\"\"'\n",
+    "    retval = textwrap.indent(retval, prefix=\" \" * col_offset)\n",
+    "    return retval"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "94a13246",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\"\"\"Sample docstring\n",
+      "\n",
+      "Args:\n",
+      "    s: sample args\n",
+      "\n",
+      "Returns:\n",
+      "    sample return\n",
+      "\"\"\"\n",
+      "\"\"\"Sample docstring\n",
+      "\n",
+      "Args:\n",
+      "    s: sample args\n",
+      "\n",
+      "Returns:\n",
+      "    sample return\n",
+      "\n",
+      "!!! note\n",
+      "\n",
+      "    The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "\"\"\"\n",
+      "    \"\"\"Sample docstring\n",
+      "\n",
+      "    Args:\n",
+      "        s: sample args\n",
+      "\n",
+      "    Returns:\n",
+      "        sample return\n",
+      "    \"\"\"\n",
+      "    \"\"\"Sample docstring\n",
+      "\n",
+      "    Args:\n",
+      "        s: sample args\n",
+      "\n",
+      "    Returns:\n",
+      "        sample return\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "        \"\"\"Sample docstring\n",
+      "\n",
+      "        Args:\n",
+      "            s: sample args\n",
+      "\n",
+      "        Returns:\n",
+      "            sample return\n",
+      "        \"\"\"\n",
+      "        \"\"\"Sample docstring\n",
+      "\n",
+      "        Args:\n",
+      "            s: sample args\n",
+      "\n",
+      "        Returns:\n",
+      "            sample return\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n"
+     ]
+    }
+   ],
+   "source": [
+    "docstring = \"\"\"Sample docstring\n",
+    "\n",
+    "    Args:\n",
+    "        s: sample args\n",
+    "\n",
+    "    Returns:\n",
+    "        sample return\n",
+    "\"\"\"\n",
+    "\n",
+    "expected_false = '''\"\"\"Sample docstring\n",
+    "\n",
+    "Args:\n",
+    "    s: sample args\n",
+    "\n",
+    "Returns:\n",
+    "    sample return\n",
+    "\"\"\"'''\n",
+    "expected_true = '''\"\"\"Sample docstring\n",
+    "\n",
+    "Args:\n",
+    "    s: sample args\n",
+    "\n",
+    "Returns:\n",
+    "    sample return\n",
+    "''' + AUTO_GEN_TXT + '\\n\"\"\"'\n",
+    "\n",
+    "for col_offset in [0, 4, 8]:\n",
+    "    for include_auto_gen_txt, expected in zip([False, True], [expected_false, expected_true]):\n",
+    "        expected = textwrap.indent(expected, prefix = \" \" * col_offset)\n",
+    "        actual = _fix_docstring_indent(docstring, col_offset, include_auto_gen_txt=include_auto_gen_txt)\n",
+    "        print(actual)\n",
+    "        assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aeeabeec",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "4\n"
+     ]
+    }
+   ],
+   "source": [
+    "linenos = _get_classes_and_functions(source)\n",
+    "classes_and_functions = [\n",
+    "    _get_code_from_source(source, start_line_no, end_line_no)\n",
+    "    for start_line_no, docstring_line_no, end_line_no, node_offset in linenos\n",
+    "]\n",
+    "print(len(classes_and_functions))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "25bb23e5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['Drive a car\\n\\n    Args:\\n        name: The name of the car\\n    ', 'Drive a car\\n\\n    Args:\\n        name: The name of the car\\n    ', 'Drive a car\\n\\n    Args:\\n        name: The name of the car\\n    ', 'Drive a car\\n\\n    Args:\\n        name: The name of the car\\n    ']\n",
+      "****************************************************************************************************\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "****************************************************************************************************\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
+      "****************************************************************************************************\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
+      "****************************************************************************************************\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n"
+     ]
+    }
+   ],
+   "source": [
+    "line_offset = 0\n",
+    "with mock_openai_create():\n",
+    "    docstrings = [\n",
+    "        _generate_docstring_using_codex(\n",
+    "            i,\n",
+    "            DEFAULT_PROMPT_TEMPLATE,\n",
+    "            model=\"code-davinci-002\",\n",
+    "            temperature=0.2,\n",
+    "            max_tokens=250,\n",
+    "            top_p=1.0,\n",
+    "            frequency_penalty=0.0,\n",
+    "            presence_penalty=0.0,\n",
+    "            stop=['\"\"\"'],\n",
+    "            n=3,\n",
+    "        )\n",
+    "        for i in classes_and_functions\n",
+    "    ]\n",
+    "\n",
+    "    print(docstrings)\n",
+    "\n",
+    "source_lines = source.split(\"\\n\")\n",
+    "offsets = [node_offset for i, _, _, node_offset in linenos]\n",
+    "\n",
+    "indented_docstrings = [\n",
+    "    _fix_docstring_indent(docstring, offset, include_auto_gen_txt=include_auto_gen_txt)\n",
+    "    for docstring, offset in zip(docstrings, offsets)\n",
+    "]\n",
+    "for i in indented_docstrings:\n",
+    "    print(\"*\" * 100)\n",
+    "    print(i)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e5c1812",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
     "\n",
-    "source = \"\"\"\n",
-    "@foo(x=5)\n",
-    "async def drive(self):\n",
-    "    print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "lineno = 2\n",
-    "expected = 3\n",
-    "actual = _get_start_line_for_class_or_func(source, lineno)\n",
-    "print(actual)\n",
     "\n",
-    "assert actual == expected\n",
+    "def _inject_docstring_to_source(\n",
+    "    source: str,\n",
+    "    indented_docstrings: List[str],\n",
+    "    linenos: List[Tuple[int, int, int, int]],\n",
+    ") -> str:\n",
+    "    \"\"\"Injects docstrings into source code\n",
     "\n",
-    "source = \"\"\"\n",
-    "@foo(x=5)\n",
-    "@bar(x=5)\n",
-    "@zar(x=5)\n",
-    "class A:\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        self.a = a\n",
-    "        \n",
-    "    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "lineno = 2\n",
-    "expected = 5\n",
-    "actual = _get_start_line_for_class_or_func(source, lineno)\n",
-    "print(actual)\n",
+    "    Args:\n",
+    "        source: The source code\n",
+    "        indented_docstrings: The docstrings to be injected\n",
+    "        linenos: The line numbers where the docstrings should be injected\n",
     "\n",
-    "assert actual == expected\n",
+    "    Returns:\n",
+    "        The source code with the injected docstrings\n",
     "\n",
-    "source = \"\"\"\n",
-    "@foo(x=5)\n",
-    "@bar(x=5)\n",
+    "    !!! note\n",
     "\n",
-    "\"\"\"\n",
-    "lineno = 2\n",
-    "expected = 2\n",
-    "actual = _get_start_line_for_class_or_func(source, lineno)\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    source_lines = source.split(\"\\n\")\n",
+    "    line_offset = 0\n",
+    "    for docstring, (_, docstring_line_no, _, _) in zip(indented_docstrings, linenos):\n",
+    "        docstring_lines = docstring.split(\"\\n\")\n",
+    "        source_lines = (\n",
+    "            source_lines[: docstring_line_no + line_offset]\n",
+    "            + docstring_lines\n",
+    "            + source_lines[docstring_line_no + line_offset :]\n",
+    "        )\n",
+    "        line_offset += len(docstring_lines)\n",
+    "    \n",
+    "    return \"\\n\".join(source_lines)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a21d302a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "class Test:\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "    CONST_VAL = 1\n",
+      "    \n",
+      "    def __init__(self, a):\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
+      "        self.a = a\n",
+      "        \n",
+      "    async def drive(self):\n",
+      "    \n",
+      "    \n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
+      "        print(f'The {self.model} is now driving.')\n",
+      "        \n",
+      "    def func_with_docstring():\n",
+      "        \"\"\" Sample docstring\"\"\"\n",
+      "        pass\n",
+      "        \n",
+      "def _check_and_add_docstrings_to_source(\n",
+      "    source: str, include_auto_gen_txt: bool, **kwargs\n",
+      ") -> str:\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "    source = _remove_auto_generated_docstring(source)    \n",
+      "    tree = ast.parse(source)\n",
+      "    line_offset = 0\n",
+      "\n",
+      "    for node in tree.body:\n",
+      "        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "            continue\n",
+      "        \n",
+      "        if ast.get_docstring(node) is not None:\n",
+      "            continue\n",
+      "\n",
+      "        source, line_offset = _add_docstring(\n",
+      "            source, node, line_offset, include_auto_gen_txt, **kwargs\n",
+      "        )\n",
+      "        if not isinstance(node, ast.ClassDef):\n",
+      "            continue\n",
+      "        # Is a class and we need to check the functions inside\n",
+      "        # 29 - 36 make it as a recursive function\n",
+      "        for f in node.body:\n",
+      "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+      "                continue\n",
+      "            \n",
+      "            if ast.get_docstring(f) is not None:\n",
+      "                continue\n",
+      "\n",
+      "            # should be a function inside the class for which there is no docstring\n",
+      "            source, line_offset = _add_docstring(\n",
+      "                source, f, line_offset, include_auto_gen_txt, **kwargs\n",
+      "            )\n",
+      "\n",
+      "    return source\n",
+      "    \n",
+      "def test_callback(tree, source, start_lineno, end_lineno):\n",
+      "    \"\"\"Prints the source code of the given node\n",
+      "\n",
+      "    Args:\n",
+      "        tree: The AST tree\n",
+      "        source: The source code\n",
+      "        start_lineno: The start line number\n",
+      "        end_lineno: The end line number\n",
+      "    \"\"\"\n",
+      "\n",
+      "    pass\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "actual = _inject_docstring_to_source(source, indented_docstrings, linenos)\n",
     "print(actual)\n",
     "\n",
-    "assert actual == expected\n",
-    "\n"
+    "docstring = '''Drive a car\n",
+    "\n",
+    "    Args:\n",
+    "        name: The name of the car\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "'''\n",
+    "\n",
+    "assert docstring in actual"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7cc5fd4f",
+   "id": "c94f1355",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
     "\n",
-    "def _add_docstring(\n",
-    "    source: str,\n",
-    "    node: Union[ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef],\n",
-    "    line_offset: int,\n",
-    ") -> Tuple[str, int]:\n",
-    "    \"\"\"Add a docstring to the given node and update the source code.\n",
+    "def _remove_auto_generated_docstring(source: str) -> str:\n",
+    "    \"\"\"Remove the autogenerated docstrings from the source code.\n",
     "\n",
     "    Args:\n",
-    "        source: the source code from the notebook cell\n",
-    "        node: the AST node representing a class definition, function definition,\n",
-    "            or async function definition\n",
-    "        line_offset: the number of lines added before the current\n",
-    "            node in the source\n",
+    "        source: The source code\n",
     "\n",
     "    Returns:\n",
-    "        A tuple containing the updated source code and the new line number offset\n",
+    "        The source code without the auto generated docstring\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
     "    \"\"\"\n",
-    "    line_no = node.lineno + line_offset\n",
-    "    \n",
-    "    # Fix for ast's node.lineno giving line number of decorator \n",
-    "    # instead of function/class definition in Python 3.7\n",
-    "    line_no = _get_start_line_for_class_or_func(source, line_no)\n",
-    "    \n",
-    "    if hasattr(node, 'end_lineno') and node.end_lineno is not None:\n",
-    "        end_line_no = node.end_lineno + line_offset\n",
-    "    else:\n",
-    "        end_line_no = _calculate_end_lineno(source, line_no)\n",
+    "    placeholder = \"{DOCSTRING_PLACEHOLDER}\"\n",
+    "    retval = re.sub(\n",
+    "        f'\"\"\"((?!\"\"\").)*?({AUTO_GEN_BODY}).*?\"\"\"', placeholder, source, flags=re.DOTALL\n",
+    "    )\n",
+    "    retval = \"\\n\".join([l for l in retval.split(\"\\n\") if l.strip() != placeholder])\n",
+    "    return retval"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b027a698",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "def decorator1(func):\n",
+      "    \"\"\"Decorator function that takes a function as an argument and returns a function.\"\"\"\n",
+      "    pass\n",
+      "    \n",
+      "def decorator2(func):\n",
+      "    pass\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "_source = (\n",
+    "    '''\n",
+    "def decorator1(func):\n",
+    "    \"\"\"Decorator function that takes a function as an argument and returns a function.\"\"\"\n",
+    "    pass\n",
     "    \n",
-    "    code = _get_code_from_source(source, line_no, end_line_no)\n",
-    "    docstring = _generate_docstring_using_codex(code)\n",
+    "def decorator2(func):\n",
+    "    \"\"\"Sample docstring\n",
+    "\n",
+    "    Args:\n",
+    "        arg 1: arg 1 description\n",
+    "        arg 2: arg 2 description\n",
+    "\n",
+    "    '''\n",
+    "    + AUTO_GEN_TXT\n",
+    "    + '''\n",
+    "    \"\"\"\n",
+    "    pass\n",
+    "'''\n",
+    ")\n",
+    "\n",
+    "expected = '''\n",
+    "def decorator1(func):\n",
+    "    \"\"\"Decorator function that takes a function as an argument and returns a function.\"\"\"\n",
+    "    pass\n",
     "    \n",
-    "    source = _inject_docstring_to_source(source, docstring, line_no, node.col_offset)\n",
-    "    line_offset += len(docstring.split(\"\\n\"))\n",
-    "    return source, line_offset"
+    "def decorator2(func):\n",
+    "    pass\n",
+    "'''\n",
+    "\n",
+    "actual = _remove_auto_generated_docstring(_source)\n",
+    "print(actual)\n",
+    "\n",
+    "assert actual == expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8ad75de7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | export\n",
+    "\n",
+    "\n",
+    "def _check_and_add_docstrings_to_source(\n",
+    "    source: str,\n",
+    "    include_auto_gen_txt: bool,\n",
+    "    recreate_auto_gen_docs: bool,\n",
+    "    **kwargs: Union[int, float, str, List[str]]\n",
+    ") -> str:\n",
+    "    \"\"\"Check and add docstrings to classes and functions that don't have one.\n",
+    "\n",
+    "    Args:\n",
+    "        source: Source code\n",
+    "        include_auto_gen_txt: Include auto gen text\n",
+    "        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\n",
+    "        kwargs: Keyword arguments\n",
+    "\n",
+    "    Returns:\n",
+    "        The source code with docstrings\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    if recreate_auto_gen_docs:\n",
+    "        source = _remove_auto_generated_docstring(source)\n",
+    "\n",
+    "    linenos = _get_classes_and_functions(source)\n",
+    "    if len(linenos) != 0:\n",
+    "        classes_and_functions = [\n",
+    "            _get_code_from_source(source, start_line_no, end_line_no)\n",
+    "            for start_line_no, docstring_line_no, end_line_no, node_offset in linenos\n",
+    "        ]\n",
+    "\n",
+    "        docstrings = [\n",
+    "            _generate_docstring_using_codex(i, DEFAULT_PROMPT_TEMPLATE, **kwargs)\n",
+    "            for i in classes_and_functions\n",
+    "        ]\n",
+    "        offsets = [node_offset for i, _, _, node_offset in linenos]\n",
+    "\n",
+    "        indented_docstrings = [\n",
+    "            _fix_docstring_indent(\n",
+    "                docstring, offset, include_auto_gen_txt=include_auto_gen_txt\n",
+    "            )\n",
+    "            for docstring, offset in zip(docstrings, offsets)\n",
+    "        ]\n",
+    "        source = _inject_docstring_to_source(source, indented_docstrings, linenos)\n",
+    "\n",
+    "    return source"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d2ffe631",
+   "id": "d4a4cf1e",
    "metadata": {},
    "outputs": [
     {
@@ -595,259 +1719,451 @@
      "output_type": "stream",
      "text": [
       "\n",
-      "class test:\n",
-      "    \"\"\"Sample docstring\n",
+      "class Test:\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    CONST_VAL = 1\n",
-      "    def __init__(self, a):\n",
-      "        \"\"\"Sample docstring\n",
       "    \n",
+      "    def __init__(self, a):\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
       "        Args:\n",
-      "            s: sample args\n",
-      "    \n",
-      "        Returns:\n",
-      "            sample return\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "        \"\"\"\n",
       "        self.a = a\n",
       "        \n",
-      "    async def drive(self):\n",
-      "        \"\"\"Sample docstring\n",
-      "    \n",
-      "        Args:\n",
-      "            s: sample args\n",
+      "def test_callback(tree, source, start_lineno, end_lineno):\n",
+      "\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "    pass\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "nl = \"\\n\"\n",
+    "_source = (\n",
+    "    \"\"\"\n",
+    "class Test:\n",
+    "    CONST_VAL = 1\n",
+    "    \n",
+    "    def __init__(self, a):\n",
+    "        self.a = a\n",
+    "        \n",
+    "def test_callback(tree, source, start_lineno, end_lineno):\n",
+    "    \\\"\"\"Check and add docstrings to source\n",
+    "\n",
+    "    Args:\n",
+    "        source: Source code\n",
+    "        include_auto_gen_txt: Include auto gen text\n",
+    "        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\n",
+    "        kwargs: Keyword arguments\n",
+    "\n",
+    "    Returns:\n",
+    "        The source code with docstrings\n",
+    "\n",
+    "    \"\"\"\n",
+    "    + f'{nl.join([textwrap.indent(l, \" \"*4) for l in AUTO_GEN_TXT.split(nl)])}'\n",
+    "    + \"\"\"\n",
+    "    \\\"\"\"\n",
+    "\n",
+    "    pass\n",
+    "\n",
+    "\"\"\"\n",
+    ")\n",
+    "with mock_openai_create():\n",
+    "    updated_source = _check_and_add_docstrings_to_source(\n",
+    "        _source,\n",
+    "        include_auto_gen_txt=True,\n",
+    "        recreate_auto_gen_docs=True,\n",
+    "        model=\"code-davinci-002\",\n",
+    "        temperature=0,\n",
+    "        max_tokens=250,\n",
+    "        top_p=1.0,\n",
+    "        frequency_penalty=0.0,\n",
+    "        presence_penalty=0.0,\n",
+    "        stop=[\"#\", '\"\"\"'],\n",
+    "    )\n",
+    "\n",
+    "ast.parse(updated_source)\n",
+    "\n",
+    "assert \"Check and add docstrings to source\" not in updated_source\n",
+    "print(updated_source)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33ac2afd",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "class Test:\n",
+      "    \"\"\"Drive a car\n",
+      "\n",
+      "    Args:\n",
+      "        name: The name of the car\n",
+      "\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "    CONST_VAL = 1\n",
       "    \n",
-      "        Returns:\n",
-      "            sample return\n",
+      "    def __init__(self, a):\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "        \"\"\"\n",
-      "        print(f'The {self.model} is now driving.')\n",
+      "        self.a = a\n",
+      "        \n",
+      "def test_callback(tree, source, start_lineno, end_lineno):\n",
+      "    \"\"\"Check and add docstrings to source\n",
+      "\n",
+      "    Args:\n",
+      "        source: Source code\n",
+      "        include_auto_gen_txt: Include auto gen text\n",
+      "        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\n",
+      "        kwargs: Keyword arguments\n",
+      "\n",
+      "    Returns:\n",
+      "        The source code with docstrings\n",
+      "\n",
+      "    \n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "    \"\"\"\n",
+      "\n",
+      "    pass\n",
+      "\n",
       "\n"
      ]
     }
    ],
    "source": [
-    "source = \"\"\"\n",
-    "class test:\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        self.a = a\n",
-    "        \n",
-    "    async def drive(self):\n",
-    "        print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "\n",
-    "expected = \"\"\"\n",
-    "class test:\n",
-    "    \\\"\"\"Sample docstring\n",
-    "\n",
-    "    Args:\n",
-    "        s: sample args\n",
-    "\n",
-    "    Returns:\n",
-    "        sample return\n",
-    "    \\\"\"\"\n",
-    "    CONST_VAL = 1\n",
-    "    def __init__(self, a):\n",
-    "        \\\"\"\"Sample docstring\n",
-    "    \n",
-    "        Args:\n",
-    "            s: sample args\n",
-    "    \n",
-    "        Returns:\n",
-    "            sample return\n",
-    "        \\\"\"\"\n",
-    "        self.a = a\n",
-    "        \n",
-    "    async def drive(self):\n",
-    "        \\\"\"\"Sample docstring\n",
-    "    \n",
-    "        Args:\n",
-    "            s: sample args\n",
-    "    \n",
-    "        Returns:\n",
-    "            sample return\n",
-    "        \\\"\"\"\n",
-    "        print(f'The {self.model} is now driving.')\n",
-    "\"\"\"\n",
-    "\n",
-    "tree = ast.parse(source)\n",
-    "line_offset = 0\n",
-    "for node in tree.body:\n",
-    "    source, line_offset = _add_docstring(source, node, line_offset)\n",
+    "with mock_openai_create():\n",
+    "    updated_source = _check_and_add_docstrings_to_source(\n",
+    "        _source,\n",
+    "        include_auto_gen_txt=True,\n",
+    "        recreate_auto_gen_docs=False,\n",
+    "        model=\"code-davinci-002\",\n",
+    "        temperature=0,\n",
+    "        max_tokens=250,\n",
+    "        top_p=1.0,\n",
+    "        frequency_penalty=0.0,\n",
+    "        presence_penalty=0.0,\n",
+    "        stop=[\"#\", '\"\"\"'],\n",
+    "    )\n",
     "\n",
-    "    for f in node.body:\n",
-    "        if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
-    "            continue\n",
-    "        source, line_offset = _add_docstring(source, f, line_offset)\n",
+    "ast.parse(updated_source)\n",
     "\n",
-    "print(source)\n",
-    "assert source == expected"
+    "assert \"Check and add docstrings to source\" in updated_source\n",
+    "print(updated_source)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "52938a28",
+   "id": "0be19a1b",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
     "\n",
-    "def _check_and_add_docstrings_to_source(source: str) -> str:\n",
-    "    \"\"\"Check for missing docstrings in the source code and add them if necessary.\n",
+    "def _get_files(p: Path) -> List[Path]:\n",
+    "    \"\"\"Get Jupyter notebooks and Python files path in the directory.\n",
     "\n",
     "    Args:\n",
-    "        source: the source code from the notebook cell\n",
+    "        p: Path to the directory\n",
     "\n",
     "    Returns:\n",
-    "        The updated source code with added docstrings\n",
+    "        A list of paths to the files in the directory\n",
+    "\n",
+    "    Raises:\n",
+    "        ValueError: If the directory does not contain any Python files or notebooks\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
     "    \"\"\"\n",
+    "    exts = [\".ipynb\", \".py\"]\n",
+    "    files = [\n",
+    "        f\n",
+    "        for f in p.rglob(\"*\")\n",
+    "        if f.suffix in exts\n",
+    "        and not any(p.startswith(\".\") for p in f.parts)\n",
+    "        and not f.name.startswith(\"_\")\n",
+    "    ]\n",
+    "\n",
+    "    if len(files) == 0:\n",
+    "        raise ValueError(\n",
+    "            f\"The directory {p.resolve()} does not contain any Python files or notebooks\"\n",
+    "        )\n",
+    "\n",
+    "    return files"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "464c1bc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/.hidden\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/test_1.ipynb\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/settings.ini\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/test.ipynb\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/_test.ipynb\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/.hidden/test_1.ipynb\n",
+      "/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/.hidden/test.ipynb\n",
+      "\n",
+      "\n",
+      "[PosixPath('/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/test_1.ipynb'), PosixPath('/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpv3nfiokv/nbs/test.ipynb')]\n"
+     ]
+    }
+   ],
+   "source": [
+    "with TemporaryDirectory() as d:\n",
+    "    nbs_path = Path(d) / \"nbs\"\n",
+    "    nbs_path.mkdir(parents=True)\n",
     "\n",
-    "    tree = ast.parse(source)\n",
-    "    line_offset = 0\n",
+    "    hidden_dir = nbs_path / \".hidden\"\n",
+    "    hidden_dir.mkdir(parents=True)\n",
     "\n",
-    "    for node in tree.body:\n",
-    "        if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
-    "            continue\n",
-    "        if ast.get_docstring(node) is not None:\n",
-    "            continue\n",
+    "    shutil.copyfile(Path(\"..\") / \"settings.ini\", nbs_path / \"settings.ini\")\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nbs_path / \"_test.ipynb\"\n",
+    "    )\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nbs_path / \"test.ipynb\"\n",
+    "    )\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nbs_path / \"test_1.ipynb\"\n",
+    "    )\n",
     "\n",
-    "        # A class or a function without docstring\n",
-    "        source, line_offset = _add_docstring(source, node, line_offset)\n",
-    "        if not isinstance(node, ast.ClassDef):\n",
-    "            continue\n",
-    "        # Is a class and we need to check the functions inside\n",
-    "        # 29 - 36 make it as a recursive function\n",
-    "        for f in node.body:\n",
-    "            if not isinstance(f, (ast.FunctionDef, ast.AsyncFunctionDef)):\n",
-    "                continue\n",
-    "            if ast.get_docstring(f) is not None:\n",
-    "                continue\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", hidden_dir / \"test.ipynb\"\n",
+    "    )\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", hidden_dir / \"test_1.ipynb\"\n",
+    "    )\n",
     "\n",
-    "            # should be a function inside the class for which there is no docstring\n",
-    "            source, line_offset = _add_docstring(source, f, line_offset)\n",
+    "    for f in nbs_path.rglob(\"*\"):\n",
+    "        print(f)\n",
     "\n",
-    "    return source"
+    "    files = _get_files(nbs_path)\n",
+    "\n",
+    "    assert len(files) == 2\n",
+    "    print(f\"\\n\\n{files}\")\n",
+    "    assert files == [nbs_path / \"test_1.ipynb\", nbs_path / \"test.ipynb\"]"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e54920d5",
+   "id": "d5924c7c",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\n",
-      "def decorator1(func):\n",
-      "    \"\"\"Sample docstring\n",
-      "\n",
-      "    Args:\n",
-      "        s: sample args\n",
-      "\n",
-      "    Returns:\n",
-      "        sample return\n",
-      "    \"\"\"\n",
-      "    def inner():\n",
-      "        func()\n",
-      "    return inner\n",
-      "\n",
-      "def decorator2(func):\n",
-      "    \"\"\"Sample docstring\n",
-      "\n",
-      "    Args:\n",
-      "        s: sample args\n",
-      "\n",
-      "    Returns:\n",
-      "        sample return\n",
-      "    \"\"\"\n",
-      "    def inner():\n",
-      "        func()\n",
-      "    return inner\n",
-      "\n",
-      "@decorator1\n",
-      "@decorator2\n",
-      "def outer_func():\n",
-      "    \"\"\"Sample docstring\n",
-      "\n",
-      "    Args:\n",
-      "        s: sample args\n",
-      "\n",
-      "    Returns:\n",
-      "        sample return\n",
-      "    \"\"\"\n",
-      "    def inner_func():\n",
-      "        print(\"Hello, World!\")\n",
-      "    inner_func()\n",
-      "\n"
+      "The directory /private/var/folders/6n/3rjds7v52cd83wqkd565db0h0000gn/T/tmpbqtl37th/nbs does not contain any Python files or notebooks\n"
      ]
     }
    ],
    "source": [
-    "source = \"\"\"\n",
-    "def decorator1(func):\n",
-    "    def inner():\n",
-    "        func()\n",
-    "    return inner\n",
+    "with pytest.raises(ValueError) as e:\n",
     "\n",
-    "def decorator2(func):\n",
-    "    def inner():\n",
-    "        func()\n",
-    "    return inner\n",
-    "\n",
-    "@decorator1\n",
-    "@decorator2\n",
-    "def outer_func():\n",
-    "    def inner_func():\n",
-    "        print(\"Hello, World!\")\n",
-    "    inner_func()\n",
-    "\"\"\"\n",
-    "actual = _check_and_add_docstrings_to_source(source)\n",
-    "print(actual)"
+    "    with TemporaryDirectory() as d:\n",
+    "        nbs_path = Path(d) / \"nbs\"\n",
+    "        nbs_path.mkdir(parents=True)\n",
+    "\n",
+    "        _get_files(nbs_path)\n",
+    "\n",
+    "print(e.value)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ba209b63",
+   "id": "93d13975",
    "metadata": {},
    "outputs": [],
    "source": [
     "# | export\n",
     "\n",
     "\n",
-    "def add_docstring_to_notebook(nb_path: Union[str, Path], version: int = 4):\n",
-    "    \"\"\"Add docstrings to the source\n",
-    "    \n",
-    "    This function reads through a Jupyter notebook cell by cell and \n",
-    "    adds docstrings for classes and methods that do not have them.\n",
-    "    \n",
+    "def _add_docstring_to_nb(\n",
+    "    file: Path,\n",
+    "    include_auto_gen_txt: bool,\n",
+    "    recreate_auto_gen_docs: bool,\n",
+    "    version: int = 4,\n",
+    "    **kwargs: Union[int, float, str, List[str]]\n",
+    ") -> None:\n",
+    "    \"\"\"Add docstrings to a Jupyter notebook.\n",
+    "\n",
     "    Args:\n",
-    "        nb_path: The notebook file path\n",
-    "        version: The version of the Jupyter notebook format\n",
-    "    \"\"\"\n",
-    "    nb_path = Path(nb_path)\n",
-    "    nb = nbformat.read(nb_path, as_version=version)\n",
+    "        file (Path): Path to the notebook file.\n",
+    "        include_auto_gen_txt (bool): If True, include a text indicating that the docstring was auto-generated.\n",
+    "        recreate_auto_gen_docs (bool): If True, recreate the docstrings even if they already exist.\n",
+    "        version (int): The version of the notebook.\n",
+    "        **kwargs: Additional arguments to be passed to the function that generates the docstrings.\n",
+    "\n",
+    "    Returns:\n",
+    "        None\n",
+    "\n",
+    "    !!! note\n",
     "\n",
-    "    for cell in nb.cells:\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    _f = nbformat.read(file, as_version=version)\n",
+    "    for cell in _f.cells:\n",
     "        if cell.cell_type == \"code\":\n",
-    "            cell[\"source\"] = _check_and_add_docstrings_to_source(cell[\"source\"])\n",
+    "            cell[\"source\"] = _check_and_add_docstrings_to_source(\n",
+    "                cell[\"source\"], include_auto_gen_txt, recreate_auto_gen_docs, **kwargs\n",
+    "            )\n",
+    "    nbformat.write(_f, file)\n",
+    "\n",
+    "\n",
+    "def _add_docstring_to_py(\n",
+    "    file: Path,\n",
+    "    include_auto_gen_txt: bool,\n",
+    "    recreate_auto_gen_docs: bool,\n",
+    "    **kwargs: Union[int, float, str, List[str]]\n",
+    ") -> None:\n",
+    "    \"\"\"Add docstrings to a python file.\n",
+    "\n",
+    "    Args:\n",
+    "        file: The path to the python file.\n",
+    "        include_auto_gen_txt: Whether to include a text indicating that the docstring is auto-generated.\n",
+    "        recreate_auto_gen_docs: Whether to recreate the docstrings even if they already exist.\n",
+    "        kwargs: Additional keyword arguments.\n",
+    "\n",
+    "    Returns:\n",
+    "        None\n",
+    "\n",
+    "    Raises:\n",
+    "        ValueError: If file is not a path to a python file.\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    with file.open(\"r\") as f:\n",
+    "        source = f.read()\n",
+    "    source = _check_and_add_docstrings_to_source(\n",
+    "        source, include_auto_gen_txt, recreate_auto_gen_docs, **kwargs\n",
+    "    )\n",
+    "    with file.open(\"w\") as f:\n",
+    "        f.write(source)\n",
+    "\n",
+    "\n",
+    "def add_docstring_to_source(\n",
+    "    path: Union[str, Path],\n",
+    "    include_auto_gen_txt: bool = True,\n",
+    "    recreate_auto_gen_docs: bool = False,\n",
+    "    model: str = \"code-davinci-002\",\n",
+    "    temperature: float = 0.2,\n",
+    "    max_tokens: int = 250,\n",
+    "    top_p: float = 1.0,\n",
+    "    n: int = 3,\n",
+    ") -> None:\n",
+    "    \"\"\"Reads a Jupyter notebook or Python file, or a directory containing these files, and adds docstrings to classes and methods that do not have them.\n",
+    "\n",
+    "    Args:\n",
+    "        path: The path to the Jupyter notebook or Python file, or a directory containing these files.\n",
+    "        include_auto_gen_txt: If set to True, a note indicating that the docstring was autogenerated by docstring-gen library will be added to the end.\n",
+    "        recreate_auto_gen_docs: If set to True, the autogenerated docstrings from the previous runs will be replaced with the new one.\n",
+    "        model: The name of the Codex model that will be used to generate docstrings.\n",
+    "        temperature: Setting the temperature close to zero produces better results, whereas higher temperatures produce more complex, and sometimes irrelevant docstrings.\n",
+    "        max_tokens: The maximum number of tokens to be used when generating a docstring for a function or class. Please note that a higher number will deplete your token quota faster.\n",
+    "        top_p: You can also specify a top-P value from 0-1 to achieve similar results to changing the temperature. According to the Open AI documentation, it is generally recommended to change either this or the temperature but not both.\n",
+    "        n: The number of docstrings to be generated for each function or class, with the best one being added to the source code. Please note that a higher number will deplete your token quota faster.\n",
     "\n",
-    "    nbformat.write(nb, nb_path)"
+    "    Returns:\n",
+    "        None\n",
+    "\n",
+    "    !!! note\n",
+    "\n",
+    "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+    "    \"\"\"\n",
+    "    path = Path(path)\n",
+    "    files = _get_files(path) if path.is_dir() else [path]\n",
+    "    frequency_penalty = 0.0\n",
+    "    presence_penalty = 0.0\n",
+    "    stop = [\"#\", '\"\"\"']\n",
+    "\n",
+    "    for file in files:\n",
+    "        if file.suffix == \".ipynb\":\n",
+    "            _add_docstring_to_nb(\n",
+    "                file=file,\n",
+    "                include_auto_gen_txt=include_auto_gen_txt,\n",
+    "                recreate_auto_gen_docs=recreate_auto_gen_docs,\n",
+    "                model=model,\n",
+    "                temperature=temperature,\n",
+    "                max_tokens=max_tokens,\n",
+    "                top_p=top_p,\n",
+    "                frequency_penalty=frequency_penalty,\n",
+    "                presence_penalty=presence_penalty,\n",
+    "                stop=stop,\n",
+    "                n=n,\n",
+    "            )\n",
+    "        else:\n",
+    "            _add_docstring_to_py(\n",
+    "                file=file,\n",
+    "                include_auto_gen_txt=include_auto_gen_txt,\n",
+    "                recreate_auto_gen_docs=recreate_auto_gen_docs,\n",
+    "                model=model,\n",
+    "                temperature=temperature,\n",
+    "                max_tokens=max_tokens,\n",
+    "                top_p=top_p,\n",
+    "                frequency_penalty=frequency_penalty,\n",
+    "                presence_penalty=presence_penalty,\n",
+    "                stop=stop,\n",
+    "                n=n,\n",
+    "            )"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5abee4bd",
+   "id": "60674494",
    "metadata": {},
    "outputs": [
     {
@@ -870,19 +2186,29 @@
       "from tempfile import TemporaryDirectory\n",
       "\n",
       "# Title\n",
+      "_source = '''\n",
+      "def decorator1(func):\n",
+      "    \"\"\"Decorator function that takes a function as an argument and returns a function.\"\"\"\n",
+      "    pass\n",
+      "    \n",
+      "def decorator2(func):\n",
+      "    pass\n",
+      "'''\n",
+      "\n",
       "# | export\n",
       "\n",
       "# Vehicle class\n",
       "class Vehicle:\n",
-      "    \"\"\"Sample docstring\n",
+      "    # Constructor function\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
-      "    # Constructor function\n",
       "    def __init__(self, brand, model, type):\n",
       "        \"\"\"Constructor function\n",
       "        \n",
@@ -899,27 +2225,29 @@
       "    \n",
       "    # fuel_up function\n",
       "    def fuel_up(self):\n",
-      "        \"\"\"Sample docstring\n",
-      "    \n",
+      "        # comment goes here\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
       "        Args:\n",
-      "            s: sample args\n",
-      "    \n",
-      "        Returns:\n",
-      "            sample return\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "        \"\"\"\n",
-      "        # comment goes here\n",
       "        self.fuel_level = self.gas_tank_size\n",
       "        return 'Gas tank is now full.'\n",
       "\n",
       "    # drive function\n",
       "    async def drive(self):\n",
-      "        \"\"\"Sample docstring\n",
-      "    \n",
+      "        \"\"\"Drive a car\n",
+      "\n",
       "        Args:\n",
-      "            s: sample args\n",
-      "    \n",
-      "        Returns:\n",
-      "            sample return\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "        \"\"\"\n",
       "        print(f'The {self.model} is now driving.')\n",
       "v = Vehicle(\"honda\", \"city\", \"hatchback\")\n",
@@ -931,13 +2259,14 @@
       "This is a sample markdown cell. The texts goes here\n",
       "\n",
       "def say_hello(name: str) -> str:\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    return f\"Hello, {name}\"\n",
       "    \n",
@@ -946,65 +2275,97 @@
       "    return f\"I won't say hello, {name}\"\n",
       "    \n",
       "async def dont_say_hello(name: str) -> str:\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    return f\"I won't say hello, {name}\"\n",
       "    \n",
       "def decorator1(func):\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    def inner():\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
       "        func()\n",
       "    return inner\n",
       "\n",
       "def decorator2(func):\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    def inner():\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
       "        func()\n",
       "    return inner\n",
       "\n",
       "@decorator1\n",
       "@decorator2\n",
       "def outer_func():\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    def inner_func():\n",
+      "        \"\"\"Drive a car\n",
+      "\n",
+      "        Args:\n",
+      "            name: The name of the car\n",
+      "\n",
+      "        !!! note\n",
+      "\n",
+      "            The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
+      "        \"\"\"\n",
       "        print(\"Hello, World!\")\n",
       "    inner_func()\n",
       "@contextmanager\n",
       "def my_context_mgr():\n",
-      "    \"\"\"Sample docstring\n",
+      "    \"\"\"Drive a car\n",
       "\n",
       "    Args:\n",
-      "        s: sample args\n",
+      "        name: The name of the car\n",
       "\n",
-      "    Returns:\n",
-      "        sample return\n",
+      "    !!! note\n",
+      "\n",
+      "        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)\n",
       "    \"\"\"\n",
       "    try:\n",
       "        yield\n",
@@ -1019,36 +2380,41 @@
     "    nbs_path = Path(d) / \"nbs\"\n",
     "    nbs_path.mkdir(parents=True)\n",
     "\n",
-    "    nb_path = nbs_path / \"test.ipynb\"\n",
-    "    shutil.copyfile(Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nb_path)\n",
-    "    \n",
-    "    assert nb_path.exists()\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nbs_path / \"test.ipynb\"\n",
+    "    )\n",
+    "    shutil.copyfile(\n",
+    "        Path(\"..\") / \"fixtures\" / \"Test_Data.ipynb\", nbs_path / \"_test.ipynb\"\n",
+    "    )\n",
+    "\n",
+    "    shutil.copyfile(Path(\"..\") / \"fixtures\" / \"test_data.py\", nbs_path / \"test_data.py\")\n",
+    "    shutil.copyfile(Path(\"..\") / \"settings.ini\", nbs_path / \"settings.ini\")\n",
     "\n",
-    "    add_docstring_to_notebook(nb_path)\n",
+    "    with mock_openai_create():\n",
+    "        add_docstring_to_source(nbs_path, recreate_auto_gen_docs=True)\n",
     "\n",
-    "    with nb_path.open(\"r\") as f:\n",
+    "    with (nbs_path / \"test.ipynb\").open(\"r\") as f:\n",
     "        nb = nbformat.read(f, as_version=4)\n",
     "\n",
     "for cell in nb.cells:\n",
-    "#     print(cell[\"source\"])\n",
-    "    \n",
     "    if cell.cell_type == \"code\":\n",
-    "        tree = ast.parse(cell[\"source\"])        \n",
+    "        tree = ast.parse(cell[\"source\"])\n",
     "        print(cell[\"source\"])\n",
     "        for node in tree.body:\n",
-    "            if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):\n",
+    "            if not isinstance(\n",
+    "                node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)\n",
+    "            ):\n",
     "                continue\n",
     "            assert ast.get_docstring(node)\n",
-    "            \n",
+    "\n",
     "    else:\n",
-    "        print(cell[\"source\"])\n",
-    "\n"
+    "        print(cell[\"source\"])"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "af3e621b",
+   "id": "ddf6664b",
    "metadata": {},
    "outputs": [],
    "source": []
diff --git a/settings.ini b/settings.ini
index ba757e3..f43e7fc 100644
--- a/settings.ini
+++ b/settings.ini
@@ -39,7 +39,6 @@ user = airtai
 
 ### Optional ###
 requirements = \
-    nbdev>=2.3.9 \
     nbformat>=5.7.1 \
     typer[all]>=0.7.0 \
     openai>=0.25.0
@@ -62,4 +61,4 @@ dev_requirements =     \
     openai>=0.25.0 \
     playwright>=1.28.0
 
-# console_scripts =
\ No newline at end of file
+console_scripts = docstring_gen=docstring_gen._cli:_app


Test patch:
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index b07727e..4a40a98 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -1,6 +1,9 @@
 name: CI
 on:  [workflow_dispatch, pull_request, push]
 
+env:
+  OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+
 jobs:
   static_check:
     runs-on: ubuntu-latest
diff --git a/fixtures/test_data.py b/fixtures/test_data.py
new file mode 100644
index 0000000..baf826e
--- /dev/null
+++ b/fixtures/test_data.py
@@ -0,0 +1,69 @@
+# %% ../../nbs/Utils.ipynb 1
+import os
+from typing import *
+from urllib.parse import urlparse
+from contextlib import contextmanager
+from pathlib import Path
+
+import nbdev
+from configparser import ConfigParser
+
+# %% ../../nbs/Utils.ipynb 3
+@contextmanager
+def set_cwd(cwd_path: Union[Path, str]):
+
+    cwd_path = Path(cwd_path)
+    original_cwd = os.getcwd()
+    os.chdir(cwd_path)
+
+    try:
+        nbdev.config.get_config.cache_clear()
+        yield
+    finally:
+        os.chdir(original_cwd)
+
+# %% ../../nbs/Utils.ipynb 5
+def get_value_from_config(root_path: str, config_name: str) -> str:
+    """Get the value from settings.ini file"""
+
+    settings_path = Path(root_path) / "settings.ini"
+    config = ConfigParser()
+    config.read(settings_path)
+    if not config.has_option("DEFAULT", config_name):
+        return ""
+    return config["DEFAULT"][config_name]
+
+# %% ../../nbs/Utils.ipynb 7
+def is_local_path(path):
+    # Check if the path is an absolute path
+    if os.path.isabs(path):
+        return True
+
+    # Check if the path is a URL with a scheme (e.g. http, https, ftp)
+    parsed_url = urlparse(path)
+    if parsed_url.scheme:
+        return False
+
+    # If the path is not an absolute path and does not have a URL scheme,
+    # it is assumed to be a local path
+    return True
+
+# %% ../../nbs/Utils.ipynb 9
+def add_counter_suffix_to_filename(src_path: Path):
+    """Add a counter suffix to the given file
+    Args:
+        src_path: The path to the file to rename.
+    """
+    parent_dir = src_path.parent
+    counter_suffix = (
+        max(
+            [
+                int(f.stem.split(".")[1])
+                for f in parent_dir.glob(f"{src_path.stem}.*.*")
+            ],
+            default=0,
+        )
+        + 1
+    )
+    dst_path = parent_dir / f"{src_path.stem}.{counter_suffix}{src_path.suffix}"
+    os.rename(src_path, dst_path)