156 lines
4.4 KiB
Python
156 lines
4.4 KiB
Python
from enum import StrEnum
|
|
from functools import lru_cache
|
|
from typing import Annotated, Any, Literal
|
|
|
|
from pydantic import (
|
|
BaseModel,
|
|
BeforeValidator,
|
|
Field,
|
|
field_validator,
|
|
model_validator,
|
|
)
|
|
|
|
from .settings import KEYWORDS_LIST_PATH
|
|
|
|
|
|
@lru_cache
|
|
def load_valid_keywords() -> set[str]:
|
|
"""Loads all valid keywords from the configuration file."""
|
|
if not KEYWORDS_LIST_PATH.exists():
|
|
return set()
|
|
|
|
keywords = set()
|
|
with open(KEYWORDS_LIST_PATH, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
# Each line can contain hierarchical tags separated by '|'
|
|
# e.g., 'architektur|gebäude' means both 'architektur' and 'gebäude' are valid.
|
|
parts = line.split("|")
|
|
current_path = ""
|
|
for part in parts:
|
|
if current_path:
|
|
current_path += "|" + part
|
|
else:
|
|
current_path = part
|
|
keywords.add(current_path)
|
|
return keywords
|
|
|
|
|
|
class LlmRequest(BaseModel):
|
|
model: str = "gpt-4o"
|
|
messages: list[LlmRequestMessage]
|
|
temperature: float | None = None
|
|
|
|
|
|
class LlmMessageRole(StrEnum):
|
|
system = "system"
|
|
user = "user"
|
|
|
|
|
|
class LlmRequestMessage(BaseModel):
|
|
role: LlmMessageRole
|
|
content: list[LlmMessageContent] | str
|
|
|
|
|
|
def transform_to_image_dict(v: Any) -> Any:
|
|
if isinstance(v, str):
|
|
return {"url": f"data:image/png;base64,{v}"}
|
|
return v
|
|
|
|
|
|
Base64ImageUrl = Annotated[dict[str, str], BeforeValidator(transform_to_image_dict)]
|
|
|
|
|
|
class LlmMessageType(StrEnum):
|
|
text = "text"
|
|
image_url = "image_url"
|
|
|
|
|
|
class LlmMessageContent(BaseModel):
|
|
type: LlmMessageType
|
|
text: str | None = None
|
|
image_url: Base64ImageUrl | None = None
|
|
|
|
@model_validator(mode="after")
|
|
def _type_matches_content(self):
|
|
if self.type == "text" and self.text is None:
|
|
raise ValueError("text must be set when type is 'text'")
|
|
if self.type == "image_url" and self.image_url is None:
|
|
raise ValueError("image_url must be set when type is 'image_url'")
|
|
return self
|
|
|
|
|
|
class LlmResponseGeneral(BaseModel):
|
|
text: str
|
|
|
|
|
|
class LlmResponseCaptionsAndKeywords(BaseModel):
|
|
caption: Annotated[
|
|
str,
|
|
Field(
|
|
description=(
|
|
"A brief but precise description of what the picture shows. "
|
|
"No interpretation; just what is visible"
|
|
)
|
|
),
|
|
]
|
|
tags: list[str] = Field(
|
|
description=(
|
|
"A list of keywords that properly tag the photo. "
|
|
"Allowed keywords: "
|
|
+ str(load_valid_keywords())
|
|
.replace("'", '"')
|
|
.replace("{", "[")
|
|
.replace("}", "]")
|
|
)
|
|
)
|
|
|
|
@field_validator("tags")
|
|
@classmethod
|
|
def validate_tags(cls, v: list[str]) -> list[str]:
|
|
valid_keywords = load_valid_keywords()
|
|
for tag in v:
|
|
if tag not in valid_keywords:
|
|
raise ValueError(
|
|
f"'{tag}' is not a valid keyword. Must be one from the configuration."
|
|
)
|
|
return v
|
|
|
|
|
|
class LlmResponseQuality(BaseModel):
|
|
rating: Annotated[
|
|
Literal["ok", "review", "not_ok"],
|
|
Field(
|
|
description=(
|
|
"General quality assessment of the photo. "
|
|
'Helping to decide whether to keep ("ok"), '
|
|
'to consider sorting out ("review"), '
|
|
'or to almost certainly sort out ("not_ok")'
|
|
)
|
|
),
|
|
]
|
|
reason: (
|
|
None
|
|
| Annotated[
|
|
Literal["blur", "exposure", "style"],
|
|
Field(
|
|
description=(
|
|
'Possible reasons for a lower-than-"ok" assessment. '
|
|
'"blur": Obviously mis-focussed or motion-blurred. '
|
|
'"exposure": severly under- or over-exposed. '
|
|
'"style": photo has severe style issues (beyond personal taste).'
|
|
)
|
|
),
|
|
]
|
|
)
|
|
|
|
@model_validator(mode="after")
|
|
def _not_ok_needs_reason(self):
|
|
if self.rating != "ok" and self.reason is None:
|
|
raise ValueError("lower-than-ok rating needs a reason")
|
|
elif self.rating == "ok" and self.reason is not None:
|
|
raise ValueError("ok rating needs no reason")
|
|
return self
|