[Prompt] Suggest how to install when ModuleNotFound

[Kyle Kelley]

When a user hits a ModuleNotFound Error, we should send a different kind of prompt and context. Sometimes the user makes a typo, and sometimes they don't have the package installed. genai should be able to figure out between the two.

Missing Package

Typo

Context

• Send the user's code
• Send the output of pip freeze (assuming this isn't too big) as a role: system message

Bonus Context / Optimization

As a bonus, if the pip freeze output is too big maybe we can use something like levenshtein distance to pull the closest string matches.

packages = !pip list --format=freeze
packages = [pkg.split('==')[0] for pkg in packages]

import Levenshtein

def find_closest_strings(target, string_list, n=20):
    """
    Finds the n closest strings in string_list to the target string.
    """
    # Compute the Levenshtein distance between the target string and each string in the list
    similarity_scores = [(string, Levenshtein.distance(target, string)) for string in string_list]

    # Sort the list based on the similarity scores
    similarity_scores.sort(key=lambda x: x[1])

    # Return the n closest strings
    return [x[0] for x in similarity_scores[:n]]

find_closest_strings("pndas", packages)

Outputs

['pandas',
 'conda',
 'dask',
 'anyio',
 'appdirs',
 'attrs',
 'dx',
 'fqdn',
 'fs',
 'genai',
 'geopandas',
 'idna',
 'jedi',
 'Jinja2',
 'openai',
 'parso',
 'partd',
 'patsy',
 'pip',
 'py']

Prompt suggestion

• Use !pip or %pip if an install is needed
• "Here are the packages installed on the system"