From 5181ea7b3de494347946f4c04c2147a08773a305 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Fri, 4 Oct 2024 13:23:46 -0500
Subject: [PATCH] load sqlite_vec in main thread db

---
 frigate/app.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/frigate/app.py b/frigate/app.py
index ca419e045..daebae8bb 100644
--- a/frigate/app.py
+++ b/frigate/app.py
@@ -9,6 +9,7 @@ from multiprocessing.synchronize import Event as MpEvent
 from typing import Any, Optional
 
 import psutil
+import sqlite_vec
 import uvicorn
 from peewee_migrate import Router
 from playhouse.sqlite_ext import SqliteExtDatabase
@@ -223,11 +224,8 @@ class FrigateApp:
 
     def init_embeddings_manager(self) -> None:
         if not self.config.semantic_search.enabled:
-            self.embeddings = None
             return
 
-        # Create a client for other processes to use
-        self.embeddings = EmbeddingsContext()
         embedding_process = util.Process(
             target=manage_embeddings,
             name="embeddings_manager",
@@ -266,6 +264,14 @@ class FrigateApp:
         ]
         self.db.bind(models)
 
+        if self.config.semantic_search.enabled:
+            # use existing db connection to load sqlite_vec extension
+            conn = self.db.connection()
+            conn.enable_load_extension(True)
+            sqlite_vec.load(conn)
+            conn.enable_load_extension(False)
+            logger.info(f"main connection: {self.db}")
+
     def check_db_data_migrations(self) -> None:
         # check if vacuum needs to be run
         if not os.path.exists(f"{CONFIG_DIR}/.exports"):
@@ -277,6 +283,14 @@ class FrigateApp:
 
             migrate_exports(self.config.ffmpeg, list(self.config.cameras.keys()))
 
+    def init_embeddings_client(self) -> None:
+        if not self.config.semantic_search.enabled:
+            self.embeddings = None
+            return
+
+        # Create a client for other processes to use
+        self.embeddings = EmbeddingsContext(self.db)
+
     def init_external_event_processor(self) -> None:
         self.external_event_processor = ExternalEventProcessor(self.config)
 
@@ -467,7 +481,7 @@ class FrigateApp:
         self.event_processor.start()
 
     def start_event_cleanup(self) -> None:
-        self.event_cleanup = EventCleanup(self.config, self.stop_event)
+        self.event_cleanup = EventCleanup(self.config, self.stop_event, self.db)
         self.event_cleanup.start()
 
     def start_record_cleanup(self) -> None:
@@ -591,6 +605,7 @@ class FrigateApp:
         self.init_go2rtc()
         self.bind_database()
         self.check_db_data_migrations()
+        self.init_embeddings_client()
         self.init_inter_process_communicator()
         self.init_dispatcher()
         self.start_detectors()