E2e workflow #21
@@ -12,8 +12,10 @@ from photools3.models import (
|
||||
LlmMessageType,
|
||||
LlmRequest,
|
||||
LlmRequestMessage,
|
||||
LlmResponseCaption,
|
||||
LlmResponseCaptionsAndKeywords,
|
||||
LlmResponseGeneral,
|
||||
LlmResponseKeywords,
|
||||
LlmResponseQuality,
|
||||
)
|
||||
from photools3.photo import Photo
|
||||
@@ -63,6 +65,15 @@ class LlmSession:
|
||||
api_key=self.api_key,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
self.client.close()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.close()
|
||||
|
||||
def client_chat_completion_create(self, request: LlmRequest, choice_no: int = 0):
|
||||
response = self.client.chat.completions.create(**request.model_dump())
|
||||
return response.choices[choice_no].message.content
|
||||
@@ -171,6 +182,38 @@ class LlmSessionCaptionsAndKeywords(LlmSession):
|
||||
)
|
||||
|
||||
|
||||
class LlmSessionCaption(LlmSession):
|
||||
response_model = LlmResponseCaption
|
||||
prompt = "Write a caption for the following photo"
|
||||
|
||||
def __init__(self, temperature: float = settings.llm_default_temperature):
|
||||
super().__init__(
|
||||
system_prompt=(
|
||||
"You are a helpful photographer's assistant that writes captions "
|
||||
"for photos. Answer in German.\n"
|
||||
"Output must be a JSON of the following schema:\n"
|
||||
f"{self.response_model.model_json_schema()}"
|
||||
),
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
|
||||
class LlmSessionKeywords(LlmSession):
|
||||
response_model = LlmResponseKeywords
|
||||
prompt = "Choose appropriate keywords for the following photo"
|
||||
|
||||
def __init__(self, temperature: float = settings.llm_default_temperature):
|
||||
super().__init__(
|
||||
system_prompt=(
|
||||
"You are a helpful photographer's assistant that chooses keywords "
|
||||
"for photos. Answer in German.\n"
|
||||
"Output must be a JSON of the following schema:\n"
|
||||
f"{self.response_model.model_json_schema()}"
|
||||
),
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
|
||||
class LlmSessionQuality(LlmSession):
|
||||
response_model = LlmResponseQuality
|
||||
prompt = "Assess the quality of the following photos to help me decide which ones to sort out"
|
||||
|
||||
@@ -87,12 +87,12 @@ class MetadataWriter:
|
||||
return metadata
|
||||
|
||||
def get_caption(self) -> str:
|
||||
result = self.metadata.get(XMP_KEYS_CAPTION[0], "")
|
||||
result = (self.metadata or {}).get(XMP_KEYS_CAPTION[0], "")
|
||||
assert isinstance(result, str)
|
||||
return result
|
||||
|
||||
def get_keywords(self) -> list[str]:
|
||||
result = self.metadata.get(XMP_KEYS_KEYWORDS[0], []) or []
|
||||
result = (self.metadata or {}).get(XMP_KEYS_KEYWORDS[0], []) or []
|
||||
if isinstance(result, str):
|
||||
result = [result]
|
||||
assert isinstance(result, list)
|
||||
@@ -108,8 +108,11 @@ class MetadataWriter:
|
||||
tags=metadata_tags,
|
||||
params=["-overwrite_original"],
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
except Exception as exc:
|
||||
log.error(f"Could not write metadata to file {self.sidecar_path}: {exc}")
|
||||
raise RuntimeError(
|
||||
f"Could not write metadata to file {self.sidecar_path}"
|
||||
) from exc
|
||||
|
||||
def write_caption(
|
||||
self, caption: Caption, include_prefix: bool = True, override: bool = True
|
||||
@@ -151,9 +154,27 @@ class MetadataWriter:
|
||||
keywords: Keywords | None = None,
|
||||
color: Color | None = None,
|
||||
):
|
||||
metadata_tags: XmpMetadata = {}
|
||||
if caption is not None:
|
||||
self.write_caption(caption)
|
||||
existing_no_prefix, _ = _separate_prefixed_caption(self.get_caption())
|
||||
if existing_no_prefix:
|
||||
existing_no_prefix = (
|
||||
f"{existing_no_prefix}."
|
||||
if not existing_no_prefix.endswith(".")
|
||||
else existing_no_prefix
|
||||
)
|
||||
existing_no_prefix = f"{existing_no_prefix} "
|
||||
caption = f"{existing_no_prefix}{_prefixed(caption)}"
|
||||
metadata_tags.update({key: caption for key in XMP_KEYS_CAPTION})
|
||||
|
||||
if keywords is not None:
|
||||
self.write_keywords(keywords)
|
||||
keywords = list(map(_prefixed, keywords))
|
||||
existing = self.get_keywords()
|
||||
keywords = list(set(existing + keywords))
|
||||
metadata_tags.update({tag: keywords for tag in XMP_KEYS_KEYWORDS})
|
||||
|
||||
if color is not None:
|
||||
self.write_color(color)
|
||||
metadata_tags["XMP-darktable:Colorlabels"] = color
|
||||
|
||||
if metadata_tags:
|
||||
self.set_metadata(metadata_tags)
|
||||
|
||||
@@ -86,6 +86,42 @@ class LlmResponseGeneral(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class LlmResponseCaption(BaseModel):
|
||||
caption: Annotated[
|
||||
str,
|
||||
Field(
|
||||
description=(
|
||||
"A brief but precise description of what the picture shows. "
|
||||
"No interpretation; just what is visible"
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class LlmResponseKeywords(BaseModel):
|
||||
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 LlmResponseCaptionsAndKeywords(BaseModel):
|
||||
caption: Annotated[
|
||||
str,
|
||||
|
||||
+13
-1
@@ -1,6 +1,7 @@
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Self
|
||||
|
||||
from PIL import Image
|
||||
|
||||
@@ -17,6 +18,15 @@ class Photo:
|
||||
self.path = path
|
||||
self.image = Image.open(path)
|
||||
|
||||
def close(self):
|
||||
self.image.close()
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.close()
|
||||
|
||||
def resize(self, image_max_size=1024) -> Photo:
|
||||
"""Limiting max dimensions while maintaining aspect ratio"""
|
||||
self.image.thumbnail((image_max_size, image_max_size), Image.Resampling.LANCZOS)
|
||||
@@ -24,7 +34,9 @@ class Photo:
|
||||
|
||||
def to_rgb(self):
|
||||
if self.image.mode != "RGB":
|
||||
self.image = self.image.convert("RGB")
|
||||
original = self.image
|
||||
self.image = original.convert("RGB")
|
||||
original.close()
|
||||
return self
|
||||
|
||||
def get_encoded(self):
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from pathlib import Path
|
||||
|
||||
from photools3.llm_session import (
|
||||
LlmSession,
|
||||
LlmSessionCaption,
|
||||
LlmSessionCaptionsAndKeywords,
|
||||
LlmSessionKeywords,
|
||||
)
|
||||
from photools3.metadata_writer import MetadataWriter
|
||||
from photools3.models import (
|
||||
LlmResponseCaption,
|
||||
LlmResponseCaptionsAndKeywords,
|
||||
LlmResponseKeywords,
|
||||
)
|
||||
from photools3.photo import Photo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
type ProcessingResponse = (
|
||||
LlmResponseCaption | LlmResponseKeywords | LlmResponseCaptionsAndKeywords
|
||||
)
|
||||
|
||||
|
||||
class ProcessingMode(StrEnum):
|
||||
KEYWORDS = "keywords"
|
||||
CAPTIONS = "captions"
|
||||
BOTH = "both"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProcessingOutcome:
|
||||
path: Path
|
||||
result: ProcessingResponse | None = None
|
||||
error: str | None = None
|
||||
|
||||
@property
|
||||
def succeeded(self) -> bool:
|
||||
return self.error is None and self.result is not None
|
||||
|
||||
|
||||
def _create_session(mode: ProcessingMode) -> LlmSession:
|
||||
if mode == ProcessingMode.KEYWORDS:
|
||||
return LlmSessionKeywords()
|
||||
if mode == ProcessingMode.CAPTIONS:
|
||||
return LlmSessionCaption()
|
||||
return LlmSessionCaptionsAndKeywords()
|
||||
|
||||
|
||||
def process_photo(
|
||||
path: Path | str,
|
||||
mode: ProcessingMode | str = ProcessingMode.BOTH,
|
||||
*,
|
||||
session: LlmSession | None = None,
|
||||
) -> ProcessingResponse:
|
||||
"""Process one photo and write the requested LLM result to its metadata.
|
||||
|
||||
If a session is supplied, its lifecycle remains the caller's responsibility.
|
||||
"""
|
||||
path = Path(path)
|
||||
mode = ProcessingMode(mode)
|
||||
owns_session = session is None
|
||||
active_session = session if session is not None else _create_session(mode)
|
||||
|
||||
try:
|
||||
with Photo(path) as photo:
|
||||
result = active_session.call([photo])
|
||||
|
||||
caption: str | None = None
|
||||
keywords: list[str] | None = None
|
||||
if mode == ProcessingMode.CAPTIONS:
|
||||
if not isinstance(result, LlmResponseCaption):
|
||||
raise TypeError("Caption mode returned an unexpected response type")
|
||||
caption = result.caption
|
||||
elif mode == ProcessingMode.KEYWORDS:
|
||||
if not isinstance(result, LlmResponseKeywords):
|
||||
raise TypeError("Keywords mode returned an unexpected response type")
|
||||
keywords = result.tags
|
||||
else:
|
||||
if not isinstance(result, LlmResponseCaptionsAndKeywords):
|
||||
raise TypeError("Both mode returned an unexpected response type")
|
||||
caption = result.caption
|
||||
keywords = result.tags
|
||||
|
||||
MetadataWriter(path).write(caption=caption, keywords=keywords)
|
||||
return result
|
||||
finally:
|
||||
if owns_session:
|
||||
active_session.close()
|
||||
|
||||
|
||||
def process_batch(
|
||||
paths: Iterable[Path | str],
|
||||
mode: ProcessingMode | str = ProcessingMode.BOTH,
|
||||
) -> list[ProcessingOutcome]:
|
||||
"""Process photos sequentially, reusing one LLM client for the whole batch.
|
||||
|
||||
A failed photo is recorded in its outcome and does not stop later photos.
|
||||
Only one decoded image is held at a time.
|
||||
"""
|
||||
mode = ProcessingMode(mode)
|
||||
session = _create_session(mode)
|
||||
outcomes: list[ProcessingOutcome] = []
|
||||
|
||||
try:
|
||||
for item in paths:
|
||||
path = Path(item)
|
||||
try:
|
||||
result = process_photo(path, mode, session=session)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to process %s", path)
|
||||
# Keep only a string, not the exception/traceback, which can
|
||||
# retain references to the just-processed image and stack.
|
||||
outcomes.append(
|
||||
ProcessingOutcome(
|
||||
path=path,
|
||||
error=f"{type(exc).__name__}: {exc}",
|
||||
)
|
||||
)
|
||||
else:
|
||||
outcomes.append(ProcessingOutcome(path=path, result=result))
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
return outcomes
|
||||
@@ -1,5 +1,6 @@
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -211,3 +212,32 @@ def test_set_keywords(make_dng_file, pic, expected_keywords):
|
||||
def test_set_keywords_on_existing():
|
||||
# TODO
|
||||
...
|
||||
|
||||
|
||||
def test_write_caption_and_keywords_in_one_metadata_update(tmp_path):
|
||||
writer = MetadataWriter(tmp_path / "photo.jpg")
|
||||
with (
|
||||
patch.object(writer, "get_caption", return_value="Original caption"),
|
||||
patch.object(writer, "get_keywords", return_value=["existing"]),
|
||||
patch.object(writer, "set_metadata") as set_metadata,
|
||||
):
|
||||
writer.write(caption="Generated caption", keywords=["architektur"])
|
||||
|
||||
set_metadata.assert_called_once()
|
||||
metadata_tags = set_metadata.call_args.args[0]
|
||||
assert metadata_tags["XMP:Description"] == (
|
||||
f"Original caption. {PREFIX}Generated caption"
|
||||
)
|
||||
assert set(metadata_tags["XMP:Subject"]) == {
|
||||
"existing",
|
||||
f"{PREFIX}architektur",
|
||||
}
|
||||
|
||||
|
||||
def test_set_metadata_surfaces_write_failures(tmp_path):
|
||||
writer = MetadataWriter(tmp_path / "photo.jpg")
|
||||
with patch("photools3.metadata_writer.exiftool.ExifToolHelper") as exif:
|
||||
exif_instance = exif.return_value.__enter__.return_value
|
||||
exif_instance.execute.side_effect = OSError("ExifTool failed")
|
||||
with pytest.raises(RuntimeError, match="Could not write metadata"):
|
||||
writer.set_metadata({"XMP:Description": "caption"})
|
||||
|
||||
@@ -48,6 +48,15 @@ def test_load_photo():
|
||||
assert isinstance(photo.image, Image.Image)
|
||||
|
||||
|
||||
def test_photo_context_manager_closes_image(tmp_path):
|
||||
path = _make_image_file(tmp_path, "RGB")
|
||||
with Photo(path) as photo:
|
||||
image = photo.image
|
||||
assert image.fp is not None
|
||||
|
||||
assert image.fp is None
|
||||
|
||||
|
||||
def test_resize_method_exists():
|
||||
"""Test that resize method exists and returns a Photo instance"""
|
||||
photo = Photo(PATHS_DNG_EXAMPLES[0])
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from photools3.models import (
|
||||
LlmResponseCaption,
|
||||
LlmResponseCaptionsAndKeywords,
|
||||
LlmResponseKeywords,
|
||||
)
|
||||
from photools3.processing import (
|
||||
ProcessingMode,
|
||||
_create_session,
|
||||
process_batch,
|
||||
process_photo,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_path(tmp_path):
|
||||
path = tmp_path / "photo.png"
|
||||
Image.new("RGB", (20, 20)).save(path)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "response", "expected_metadata"),
|
||||
[
|
||||
(
|
||||
ProcessingMode.KEYWORDS,
|
||||
LlmResponseKeywords(tags=["architektur"]),
|
||||
{"caption": None, "keywords": ["architektur"]},
|
||||
),
|
||||
(
|
||||
ProcessingMode.CAPTIONS,
|
||||
LlmResponseCaption(caption="Ein Haus"),
|
||||
{"caption": "Ein Haus", "keywords": None},
|
||||
),
|
||||
(
|
||||
ProcessingMode.BOTH,
|
||||
LlmResponseCaptionsAndKeywords(caption="Ein Haus", tags=["architektur"]),
|
||||
{"caption": "Ein Haus", "keywords": ["architektur"]},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_process_photo_writes_only_requested_metadata(
|
||||
image_path, mode, response, expected_metadata
|
||||
):
|
||||
session = MagicMock()
|
||||
session.call.return_value = response
|
||||
|
||||
with (
|
||||
patch("photools3.processing._create_session", return_value=session),
|
||||
patch("photools3.processing.MetadataWriter") as writer_class,
|
||||
):
|
||||
result = process_photo(image_path, mode)
|
||||
|
||||
assert result == response
|
||||
session.call.assert_called_once()
|
||||
session.close.assert_called_once()
|
||||
writer_class.assert_called_once_with(image_path)
|
||||
writer_class.return_value.write.assert_called_once_with(**expected_metadata)
|
||||
|
||||
|
||||
def test_process_photo_does_not_close_a_caller_owned_session(image_path):
|
||||
session = MagicMock()
|
||||
session.call.return_value = LlmResponseCaption(caption="Ein Haus")
|
||||
|
||||
with patch("photools3.processing.MetadataWriter"):
|
||||
process_photo(image_path, ProcessingMode.CAPTIONS, session=session)
|
||||
|
||||
session.close.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "constructor"),
|
||||
[
|
||||
(ProcessingMode.KEYWORDS, "LlmSessionKeywords"),
|
||||
(ProcessingMode.CAPTIONS, "LlmSessionCaption"),
|
||||
(ProcessingMode.BOTH, "LlmSessionCaptionsAndKeywords"),
|
||||
],
|
||||
)
|
||||
def test_create_session_selects_mode_specific_response_model(mode, constructor):
|
||||
with patch(f"photools3.processing.{constructor}") as session_class:
|
||||
assert _create_session(mode) is session_class.return_value
|
||||
|
||||
session_class.assert_called_once_with()
|
||||
|
||||
|
||||
def test_process_batch_is_sequential_reuses_session_and_continues_after_error():
|
||||
active_photos = 0
|
||||
max_active_photos = 0
|
||||
|
||||
class TrackingPhoto:
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
|
||||
def __enter__(self):
|
||||
nonlocal active_photos, max_active_photos
|
||||
active_photos += 1
|
||||
max_active_photos = max(max_active_photos, active_photos)
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
nonlocal active_photos
|
||||
active_photos -= 1
|
||||
|
||||
session = MagicMock()
|
||||
response = LlmResponseCaptionsAndKeywords(caption="Ein Haus", tags=["architektur"])
|
||||
session.call.side_effect = [response, RuntimeError("temporary error"), response]
|
||||
paths = [Path("one.jpg"), Path("two.jpg"), Path("three.jpg")]
|
||||
|
||||
with (
|
||||
patch("photools3.processing._create_session", return_value=session),
|
||||
patch("photools3.processing.Photo", TrackingPhoto),
|
||||
patch("photools3.processing.MetadataWriter") as writer_class,
|
||||
):
|
||||
outcomes = process_batch(paths, ProcessingMode.BOTH)
|
||||
|
||||
assert [outcome.succeeded for outcome in outcomes] == [True, False, True]
|
||||
assert [outcome.path for outcome in outcomes] == paths
|
||||
assert outcomes[1].error == "RuntimeError: temporary error"
|
||||
assert max_active_photos == 1
|
||||
assert active_photos == 0
|
||||
assert session.call.call_count == 3
|
||||
assert all(len(call.args[0]) == 1 for call in session.call.call_args_list)
|
||||
assert writer_class.return_value.write.call_count == 2
|
||||
session.close.assert_called_once()
|
||||
Reference in New Issue
Block a user