From 8ed86a907ab8c0722481d2edb3e47d4dc9564764 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Wed, 25 Sep 2024 08:53:49 -0500
Subject: [PATCH] Add ability to restrict genai to labels and zones at the
 camera level

---
 docs/docs/configuration/genai.md     |  7 ++++++-
 docs/docs/configuration/reference.md |  6 ++++++
 frigate/config.py                    | 15 +++++++++++++++
 frigate/embeddings/maintainer.py     |  8 ++++++++
 web/src/types/frigateConfig.ts       |  2 ++
 5 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/docs/docs/configuration/genai.md b/docs/docs/configuration/genai.md
index 26954effe..b871d97fd 100644
--- a/docs/docs/configuration/genai.md
+++ b/docs/docs/configuration/genai.md
@@ -128,7 +128,7 @@ genai:
     car: "Label the primary vehicle in these images with just the name of the company if it is a delivery vehicle, or the color make and model."
 ```
 
-Prompts can also be overriden at the camera level to provide a more detailed prompt to the model about your specific camera, if you desire.
+Prompts can also be overriden at the camera level to provide a more detailed prompt to the model about your specific camera, if you desire. You can also optionally specify `labels` and `required_zones` to only generate descriptions for certain labels or zones.
 
 ```yaml
 cameras:
@@ -138,6 +138,11 @@ cameras:
       object_prompts:
         person: "Describe the main person in these images (gender, age, clothing, activity, etc). Do not include where the activity is occurring (sidewalk, concrete, driveway, etc). If delivering a package, include the company the package is from."
         cat: "Describe the cat in these images (color, size, tail). Indicate whether or not the cat is by the flower pots. If the cat is chasing a mouse, make up a name for the mouse."
+      labels:
+        - person
+        - cat
+      required_zones:
+        - steps
 ```
 
 ### Experiment with prompts
diff --git a/docs/docs/configuration/reference.md b/docs/docs/configuration/reference.md
index 020e59979..27aeaff77 100644
--- a/docs/docs/configuration/reference.md
+++ b/docs/docs/configuration/reference.md
@@ -723,6 +723,12 @@ cameras:
       # Format: {label}: {prompt}
       object_prompts:
         person: "My special person prompt."
+      # Optional: labels that qualify as an alert (default: all labels that are tracked)
+      labels:
+        - person
+        - cat
+      # Optional: Restrict generation to objects that entered any of the listed zones (default: none, all zones qualify)
+      required_zones: []
 
 # Optional
 ui:
diff --git a/frigate/config.py b/frigate/config.py
index ffd5afb02..f82c71ad7 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -810,6 +810,21 @@ class GenAICameraConfig(BaseModel):
         title="Default caption prompt.",
     )
     object_prompts: Dict[str, str] = Field(default={}, title="Object specific prompts.")
+    labels: Optional[List[str]] = Field(
+        default=None, title="Labels to run generative AI for."
+    )
+    required_zones: Union[str, List[str]] = Field(
+        default_factory=list,
+        title="List of required zones to be entered in order to run generative AI.",
+    )
+
+    @field_validator("required_zones", mode="before")
+    @classmethod
+    def validate_required_zones(cls, v):
+        if isinstance(v, str) and "," not in v:
+            return [v]
+
+        return v
 
 
 class AudioConfig(FrigateBaseModel):
diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py
index eea5c0f86..a18f5518c 100644
--- a/frigate/embeddings/maintainer.py
+++ b/frigate/embeddings/maintainer.py
@@ -127,6 +127,14 @@ class EmbeddingMaintainer(threading.Thread):
                     camera_config.genai.enabled
                     and self.genai_client is not None
                     and event.data.get("description") is None
+                    and (
+                        camera_config.genai.labels is None
+                        or event.label in camera_config.genai.labels
+                    )
+                    and (
+                        camera_config.genai.required_zones is None
+                        or set(event.zones) & set(camera_config.genai.zones)
+                    )
                 ):
                     # Generate the description. Call happens in a thread since it is network bound.
                     threading.Thread(
diff --git a/web/src/types/frigateConfig.ts b/web/src/types/frigateConfig.ts
index a12c24248..2ab865c5a 100644
--- a/web/src/types/frigateConfig.ts
+++ b/web/src/types/frigateConfig.ts
@@ -306,6 +306,8 @@ export interface FrigateConfig {
     model: string;
     prompt: string;
     object_prompts: { [key: string]: string };
+    required_zones: string[];
+    labels: string[];
   };
 
   go2rtc: {