Skip to content

Embeddings API

locallore.embeddings

MAX_EMBEDDING_CHARS = 8000 module-attribute

MODEL_CHECKSUM_FILE = '.locallore-model.sha256' module-attribute

FloatArray = NDArray[np.float32] module-attribute

Embedder

Bases: Protocol

model_id property

dimension property

encode(texts)

Source code in src/locallore/embeddings.py 
24
def encode(self, texts: Sequence[str]) -> FloatArray: ...

encode_query(query)

Source code in src/locallore/embeddings.py 
26
def encode_query(self, query: str) -> FloatArray: ...

FastEmbedder(model_name, cache_dir, dimension, *, model_id=None)

Source code in src/locallore/embeddings.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def __init__(
    self,
    model_name: str,
    cache_dir: Path,
    dimension: int,
    *,
    model_id: str | None = None,
) -> None:
    from fastembed import TextEmbedding

    self._dimension = dimension
    self._model_id = model_id or embedding_model_id(model_name, cache_dir)
    self._model = TextEmbedding(
        model_name=model_name,
        cache_dir=str(cache_dir),
        local_files_only=True,
    )

model_id property

dimension property

encode(texts)

Source code in src/locallore/embeddings.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def encode(self, texts: Sequence[str]) -> FloatArray:
    if not texts:
        return np.empty((0, self.dimension), dtype=np.float32)
    vectors = np.asarray(list(self._model.embed(list(texts))), dtype=np.float32)
    if vectors.ndim != 2 or vectors.shape[1] != self.dimension:
        raise ValueError(
            f"model returned shape {vectors.shape}, expected (*, {self.dimension})"
        )
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    if np.any(norms == 0):
        raise ValueError("embedding model returned a zero-length vector")
    return cast(FloatArray, vectors / norms)

encode_query(query)

Source code in src/locallore/embeddings.py 
109
110
111
112
113
114
115
116
117
118
def encode_query(self, query: str) -> FloatArray:
    vectors = np.asarray(list(self._model.query_embed(query)), dtype=np.float32)
    if vectors.shape != (1, self.dimension):
        raise ValueError(
            f"model returned shape {vectors.shape}, expected (1, {self.dimension})"
        )
    norm = np.linalg.norm(vectors[0])
    if norm == 0:
        raise ValueError("embedding model returned a zero-length query vector")
    return cast(FloatArray, vectors[0] / norm)

encode_vector(vector)

Source code in src/locallore/embeddings.py 
29
30
def encode_vector(vector: FloatArray) -> bytes:
    return np.asarray(vector, dtype="<f4").tobytes(order="C")

decode_vector(value, dimension)

Source code in src/locallore/embeddings.py 
33
34
35
36
37
38
39
def decode_vector(value: bytes, dimension: int) -> FloatArray:
    vector = np.frombuffer(value, dtype="<f4")
    if vector.size != dimension:
        raise ValueError(
            f"stored vector has dimension {vector.size}, expected {dimension}"
        )
    return vector

embedding_model_id(model_name, cache_dir)

Source code in src/locallore/embeddings.py 
65
66
def embedding_model_id(model_name: str, cache_dir: Path) -> str:
    return f"{model_name}@sha256:{_directory_checksum(cache_dir)}"

has_pending_messages(connection, model_id, dimension)

Source code in src/locallore/embeddings.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def has_pending_messages(
    connection: sqlite3.Connection,
    model_id: str,
    dimension: int,
) -> bool:
    row = connection.execute(
        "SELECT EXISTS("
        "SELECT 1 FROM messages m "
        "LEFT JOIN embeddings e ON e.message_id = m.id "
        "WHERE m.role IN ('user', 'assistant') AND length(trim(m.text)) >= 3 "
        "AND length(m.text) <= ? "
        "AND (e.message_id IS NULL OR e.model_id != ? OR e.dimension != ? "
        "OR e.content_hash != m.content_hash)"
        ")",
        (MAX_EMBEDDING_CHARS, model_id, dimension),
    ).fetchone()
    return bool(row[0])

embed_pending_messages(connection, embedder, *, batch_size=64)

Source code in src/locallore/embeddings.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def embed_pending_messages(
    connection: sqlite3.Connection,
    embedder: Embedder,
    *,
    batch_size: int = 64,
) -> int:
    if batch_size < 1:
        raise ValueError("embedding batch size must be positive")
    rows = connection.execute(
        "SELECT m.id, m.text, m.content_hash FROM messages m "
        "LEFT JOIN embeddings e ON e.message_id = m.id "
        "WHERE m.role IN ('user', 'assistant') AND length(trim(m.text)) >= 3 "
        "AND length(m.text) <= ? "
        "AND (e.message_id IS NULL OR e.model_id != ? OR e.dimension != ? "
        "OR e.content_hash != m.content_hash) ORDER BY m.rowid",
        (MAX_EMBEDDING_CHARS, embedder.model_id, embedder.dimension),
    ).fetchall()
    embedded = 0
    for offset in range(0, len(rows), batch_size):
        batch = rows[offset : offset + batch_size]
        vectors = embedder.encode([row["text"] for row in batch])
        if vectors.shape != (len(batch), embedder.dimension):
            raise ValueError("embedder returned an unexpected batch shape")
        with connection:
            for row, vector in zip(batch, vectors, strict=True):
                connection.execute(
                    "INSERT INTO embeddings(message_id, model_id, dimension, content_hash, vector) "
                    "VALUES (?, ?, ?, ?, ?) ON CONFLICT(message_id) DO UPDATE SET "
                    "model_id=excluded.model_id, dimension=excluded.dimension, "
                    "content_hash=excluded.content_hash, vector=excluded.vector",
                    (
                        row["id"],
                        embedder.model_id,
                        embedder.dimension,
                        row["content_hash"],
                        encode_vector(vector),
                    ),
                )
        embedded += len(batch)
    return embedded