mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-26 19:06:11 +01:00
63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
|
from enum import Enum
|
||
|
from typing import Optional, Union
|
||
|
|
||
|
from pydantic import BaseModel, Field, field_validator
|
||
|
|
||
|
from ..base import FrigateBaseModel
|
||
|
from ..env import EnvString
|
||
|
|
||
|
__all__ = ["GenAIConfig", "GenAICameraConfig", "GenAIProviderEnum"]
|
||
|
|
||
|
|
||
|
class GenAIProviderEnum(str, Enum):
|
||
|
openai = "openai"
|
||
|
gemini = "gemini"
|
||
|
ollama = "ollama"
|
||
|
|
||
|
|
||
|
# uses BaseModel because some global attributes are not available at the camera level
|
||
|
class GenAICameraConfig(BaseModel):
|
||
|
enabled: bool = Field(default=False, title="Enable GenAI for camera.")
|
||
|
prompt: str = Field(
|
||
|
default="Describe the {label} in the sequence of images with as much detail as possible. Do not describe the background.",
|
||
|
title="Default caption prompt.",
|
||
|
)
|
||
|
object_prompts: dict[str, str] = Field(
|
||
|
default_factory=dict, title="Object specific prompts."
|
||
|
)
|
||
|
|
||
|
objects: Union[str, list[str]] = Field(
|
||
|
default_factory=list,
|
||
|
title="List of objects 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 GenAIConfig(FrigateBaseModel):
|
||
|
enabled: bool = Field(default=False, title="Enable GenAI.")
|
||
|
prompt: str = Field(
|
||
|
default="Describe the {label} in the sequence of images with as much detail as possible. Do not describe the background.",
|
||
|
title="Default caption prompt.",
|
||
|
)
|
||
|
object_prompts: dict[str, str] = Field(
|
||
|
default_factory=dict, title="Object specific prompts."
|
||
|
)
|
||
|
|
||
|
api_key: Optional[EnvString] = Field(default=None, title="Provider API key.")
|
||
|
base_url: Optional[str] = Field(default=None, title="Provider base url.")
|
||
|
model: str = Field(default="gpt-4o", title="GenAI model.")
|
||
|
provider: GenAIProviderEnum = Field(
|
||
|
default=GenAIProviderEnum.openai, title="GenAI provider."
|
||
|
)
|