Datasette plugin adding a llm_embed(model_id, text) SQL function
from datasette import hookimpl
import json
import llm


def llm_embed_factory(datasette):
    config = datasette.plugin_config("datasette-llm-embed") or {}
    keys = config.get("keys") or {}

    def llm_embed(model_id, text):
        try:
            model = llm.get_embedding_model(model_id)
            if model.model_id in keys:
                model.key = keys[model.model_id]
            return llm.encode(model.embed(text))
        except Exception as e:
            return str(e)

    return llm_embed


def llm_embed_cosine(a, b):
    try:
        return llm.cosine_similarity(llm.decode(a), llm.decode(b))
    except Exception as e:
        return str(e)


def llm_embed_decode(blob):
    return json.dumps(llm.decode(blob))


@hookimpl
def prepare_connection(datasette, conn):
    conn.create_function("llm_embed_decode", 1, llm_embed_decode)
    conn.create_function("llm_embed", 2, llm_embed_factory(datasette))
    conn.create_function("llm_embed_cosine", 2, llm_embed_cosine)
    conn.create_aggregate("llm_embed_average", 1, AverageVectorAgg)


class AverageVectorAgg:
    with_scores = False

    def __init__(self):
        self.accumulated = []
        self.vector_size = 0

    def step(self, embedding):
        vector = llm.decode(embedding)
        if len(self.accumulated) == 0:
            self.accumulated = list(vector)
        else:
            for i in range(len(self.accumulated)):
                self.accumulated[i] += vector[i]
        self.vector_size += 1

    def finalize(self):
        vector = [item / self.vector_size for item in self.accumulated]
        return llm.encode(vector)

datasette_llm_embed.py
simonw/datasette-llm-embed
import pytest
from llm.plugins import pm
import llm


class EmbedDemo(llm.EmbeddingModel):
    model_id = "embed-demo"

    def embed_batch(self, texts):
        for text in texts:
            words = text.split()[:16]
            embedding = [len(word) for word in words]
            # Pad with 0 up to 16 words
            embedding += [0] * (16 - len(embedding))
            yield embedding


class EmbedKeyDemo(llm.EmbeddingModel):
    model_id = "embed-key-demo"

    def embed_batch(self, texts):
        for text in texts:
            words = text.split()[:15]
            embedding = [len(word) for word in words]
            # Pad with 0 up to 15 words
            embedding += [0] * (15 - len(embedding))
            # Last word is the length of the key
            embedding.append(len(self.key))
            yield embedding


@pytest.fixture
def embed_demo():
    return EmbedDemo()


@pytest.fixture
def embed_key_demo():
    return EmbedKeyDemo()


@pytest.fixture(autouse=True)
def register_embed_demo_model(embed_demo, embed_key_demo):
    class MockModelsPlugin:
        __name__ = "MockModelsPlugin"

        @llm.hookimpl
        def register_embedding_models(self, register):
            register(embed_demo)
            register(embed_key_demo)

    pm.register(MockModelsPlugin(), name="undo-mock-models-plugin")
    try:
        yield
    finally:
        pm.unregister(name="undo-mock-models-plugin")

tests/conftest.py
simonw/datasette-llm-embed
import base64
from datasette.app import Datasette
import llm
import pytest


@pytest.mark.asyncio
async def test_llm_embed():
    ds = Datasette()
    await ds.invoke_startup()
    response = await ds.client.get(
        "/_memory.json",
        params={
            "_shape": "array",
            "sql": "select llm_embed('embed-demo', 'hello world') as e",
        },
    )
    assert response.status_code == 200
    encoded = response.json()[0]["e"]["encoded"]
    assert llm.decode(base64.b64decode(encoded)) == (
        5.0,
        5.0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    )


@pytest.mark.asyncio
async def test_llm_embed_cosine():
    ds = Datasette()
    response = await ds.client.get(
        "/_memory.json",
        params={
            "_shape": "array",
            "sql": """
            select llm_embed_cosine(
                llm_embed('embed-demo', 'hello world'),
                llm_embed('embed-demo', 'hello again')
            ) as cosine""",
        },
    )
    assert response.status_code == 200
    pytest.approx(response.json()[0]["cosine"], 0.0001) == 0.9999


@pytest.mark.asyncio
async def test_llm_embed_uses_key():
    ds = Datasette(
        metadata={
            "plugins": {
                "datasette-llm-embed": {
                    "keys": {
                        "embed-key-demo": "default-key",
                        "aliased": "alias",
                    }
                }
            }
        }
    )
    response = await ds.client.get(
        "/_memory.json",
        params={
            "_shape": "array",
            "sql": "select llm_embed('embed-key-demo', 'hello world') as e",
        },
    )
    assert response.status_code == 200
    encoded = response.json()[0]["e"]["encoded"]
    assert llm.decode(base64.b64decode(encoded)) == (
        5.0,
        5.0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        11.0,
    )


@pytest.mark.asyncio
async def test_llm_embed_decode():
    ds = Datasette()
    response = await ds.client.get(
        "/_memory.json",
        params={
            "_shape": "array",
            "_json": "decoded",
            "sql": """
            select llm_embed_decode(
                llm_embed('embed-demo', 'hello world')
            ) as decoded""",
        },
    )
    assert response.status_code == 200
    assert response.json()[0]["decoded"] == [
        5.0,
        5.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
    ]

tests/test_llm_embed.py
simonw/datasette-llm-embed