diff --git a/src/photools3/llm_session.py b/src/photools3/llm_session.py index c79a19f..14922af 100644 --- a/src/photools3/llm_session.py +++ b/src/photools3/llm_session.py @@ -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" diff --git a/src/photools3/metadata_writer.py b/src/photools3/metadata_writer.py index 7370838..4df2a87 100644 --- a/src/photools3/metadata_writer.py +++ b/src/photools3/metadata_writer.py @@ -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) diff --git a/src/photools3/models.py b/src/photools3/models.py index f3b27da..d545c00 100644 --- a/src/photools3/models.py +++ b/src/photools3/models.py @@ -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, diff --git a/src/photools3/photo.py b/src/photools3/photo.py index 738254f..f57b0ce 100644 --- a/src/photools3/photo.py +++ b/src/photools3/photo.py @@ -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): diff --git a/src/photools3/processing.py b/src/photools3/processing.py new file mode 100644 index 0000000..5e71ee1 --- /dev/null +++ b/src/photools3/processing.py @@ -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 diff --git a/test/test_metadata_writer.py b/test/test_metadata_writer.py index b2b96a6..0c8aed1 100644 --- a/test/test_metadata_writer.py +++ b/test/test_metadata_writer.py @@ -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"}) diff --git a/test/test_photo.py b/test/test_photo.py index 0c924fe..7d8a61f 100644 --- a/test/test_photo.py +++ b/test/test_photo.py @@ -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]) diff --git a/test/test_processing.py b/test/test_processing.py new file mode 100644 index 0000000..c1a27d3 --- /dev/null +++ b/test/test_processing.py @@ -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()