mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-07-23 12:19:04 +03:00
Josh HawkinsandGitHub
c0cf08ab4a
CI / AMD64 Build (push) Canceled after 0s
CI / ARM Build (push) Canceled after 0s
CI / Jetson Jetpack 6 (push) Canceled after 0s
CI / AMD64 Extra Build (push) Canceled after 0s
CI / ARM Extra Build (push) Canceled after 0s
CI / Synaptics Build (push) Canceled after 0s
CI / Assemble and push default build (push) Canceled after 0s
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
import logging
|
|
import sqlite3
|
|
from typing import Any
|
|
|
|
import regex
|
|
from playhouse.sqliteq import SqliteQueueDatabase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
REGEXP_TIMEOUT_SECONDS = 1.0
|
|
|
|
|
|
class SqliteVecQueueDatabase(SqliteQueueDatabase):
|
|
def __init__(
|
|
self, *args: Any, load_vec_extension: bool = False, **kwargs: Any
|
|
) -> None:
|
|
self.load_vec_extension: bool = load_vec_extension
|
|
# no extension necessary, sqlite will load correctly for each platform
|
|
self.sqlite_vec_path = "/usr/local/lib/vec0"
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _connect(self, *args: Any, **kwargs: Any) -> sqlite3.Connection:
|
|
conn: sqlite3.Connection = super()._connect(*args, **kwargs) # type: ignore[misc]
|
|
if self.load_vec_extension:
|
|
self._load_vec_extension(conn)
|
|
|
|
# register REGEXP support
|
|
self._register_regexp(conn)
|
|
|
|
return conn
|
|
|
|
def _load_vec_extension(self, conn: sqlite3.Connection) -> None:
|
|
conn.enable_load_extension(True)
|
|
|
|
try:
|
|
conn.load_extension(self.sqlite_vec_path)
|
|
except conn.OperationalError:
|
|
logger.error("Unable to load the sqlite-vec extension")
|
|
self.load_vec_extension = False
|
|
finally:
|
|
conn.enable_load_extension(False)
|
|
|
|
def _register_regexp(self, conn: sqlite3.Connection) -> None:
|
|
def regexp(expr: str, item: str | None) -> bool:
|
|
if item is None:
|
|
return False
|
|
try:
|
|
return (
|
|
regex.search(expr, item, timeout=REGEXP_TIMEOUT_SECONDS) is not None
|
|
)
|
|
except (regex.error, TimeoutError):
|
|
return False
|
|
|
|
conn.create_function("REGEXP", 2, regexp)
|
|
|
|
def _delete_embeddings(self, table: str, event_ids: list[str]) -> None:
|
|
"""Delete embeddings for the given events, if the table exists.
|
|
|
|
Embeddings outlive the events they belong to when semantic search is
|
|
disabled, so deletes are attempted regardless of the current config.
|
|
"""
|
|
if not event_ids or not self.load_vec_extension:
|
|
return
|
|
|
|
# the embeddings tables are only created once semantic search has run
|
|
cursor = self.execute_sql(
|
|
"SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?",
|
|
(table,),
|
|
)
|
|
|
|
if cursor.fetchone() is None:
|
|
logger.debug("Skipping %s cleanup, table does not exist", table)
|
|
return
|
|
|
|
ids = ",".join(["?" for _ in event_ids])
|
|
self.execute_sql(f"DELETE FROM {table} WHERE id IN ({ids})", event_ids)
|
|
|
|
def delete_embeddings_thumbnail(self, event_ids: list[str]) -> None:
|
|
self._delete_embeddings("vec_thumbnails", event_ids)
|
|
|
|
def delete_embeddings_description(self, event_ids: list[str]) -> None:
|
|
self._delete_embeddings("vec_descriptions", event_ids)
|
|
|
|
def drop_embeddings_tables(self) -> None:
|
|
self.execute_sql("""
|
|
DROP TABLE vec_descriptions;
|
|
""")
|
|
self.execute_sql("""
|
|
DROP TABLE vec_thumbnails;
|
|
""")
|
|
|
|
def create_embeddings_tables(self) -> None:
|
|
"""Create vec0 virtual table for embeddings"""
|
|
self.execute_sql("""
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS vec_thumbnails USING vec0(
|
|
id TEXT PRIMARY KEY,
|
|
thumbnail_embedding FLOAT[768] distance_metric=cosine
|
|
);
|
|
""")
|
|
self.execute_sql("""
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS vec_descriptions USING vec0(
|
|
id TEXT PRIMARY KEY,
|
|
description_embedding FLOAT[768] distance_metric=cosine
|
|
);
|
|
""")
|