datasette-llm-embed by simonw

27 downloads this week        Star

README source code


PyPI Changelog Tests License

Datasette plugin adding a llm_embed(model_id, text) SQL function.


datasette install datasette-llm-embed


Adds a SQL function that can be called like this:

select llm_embed('sentence-transformers/all-mpnet-base-v2', 'This is some text')

This embeds the provided text using the specified embedding model and returns a binary blob, suitable for use with plugins such as datasette-faiss.

The models need to be installed using LLM plugins such as llm-sentence-transformers.

Use llm_embed_cosine(a, b) to calculate cosine similarity between two vector blobs:

select llm_embed_cosine(
    llm_embed('sentence-transformers/all-mpnet-base-v2', 'This is some text'),
    llm_embed('sentence-transformers/all-mpnet-base-v2', 'This is some other text')

The llm_embed_decode() function can be used to decode a binary BLOB into a JSON array of floats:

select llm_embed_decode(
    llm_embed('sentence-transformers/all-mpnet-base-v2', 'This is some text')

Models that require API keys

If your embedding model needs an API key - for example the ada-002 model from OpenAI - you can configure that key in metadata.yml (or JSON) like this:

        $env: OPENAI_API_KEY

The key here should be the full model ID of the model - not an alias.

You can then set the OPENAI_API_KEY environment variable to the key you want to use before starting Datasette:

export OPENAI_API_KEY=sk-1234567890

Once configured, calls like this will use the API key that has been provided:

select llm_embed('ada-002', 'This is some text')


To set up this plugin locally, first checkout the code. Then create a new virtual environment:

cd datasette-llm-embed
python3 -m venv venv
source venv/bin/activate

Now install the dependencies and test dependencies:

pip install -e '.[test]'
To run the tests: