first successful run of photo llm (#11)

Reviewed-on: #11
Co-authored-by: Mathias Scheider <mathias2784@web.de>

src/photools3/llm_session.py

@@ -1,31 +1,112 @@
-from openai import OpenAI
+import json
-from openai.types.chat import ChatCompletion
+import re
 
-from photools3.models import LlmRequest
+import openai
+from openai import OpenAI
+
+from photools3.models import (
+    LlmMessageContent,
+    LlmMessageRole,
+    LlmMessageType,
+    LlmRequest,
+    LlmRequestMessage,
+)
+from photools3.photo import Photo
 from photools3.settings import Settings
 
-settings = Settings()
+settings = Settings()  # type: ignore[reportCallIssue]
 
 
 class LlmSession:
     base_url: str = str(settings.llm_api_base_url)
     api_key: str = settings.llm_api_key
 
-    def __init__(self, system_prompt: str):
+    def __init__(self, system_prompt: str, temperature: float = 1.0):
         self._models: list = []
 
         self.system_prompt = system_prompt
+        self.temperature = temperature
 
         self.client = OpenAI(
             base_url=self.base_url,
             api_key=self.api_key,
         )
 
+    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
+
+    def is_api_reachable(self) -> bool:
+        try:
+            _ = self.client.models.list().data
+            return True
+        except openai.APIConnectionError:
+            return False
+
     @property
     def models(self):
         if not self._models:
             self._models = self.client.models.list().data
         return self._models
 
-    def call_raw(self, request: LlmRequest) -> ChatCompletion:
+    def call(
-        return self.client.chat.completions.create(**request.model_dump())
+        self,
+        prompt: str,
+        photos: list[Photo],
+        model_id=None,
+    ):
+        if not model_id:
+            model_id = self.models[0].id
+
+        request = self._build_request(prompt, photos, model_id)
+        result_raw = self.client_chat_completion_create(request)
+
+        try:
+            result = json.loads(result_raw)
+        except json.JSONDecodeError as e:
+            try:
+                json_str = re.sub(
+                    r"^```json(.*)```$", r"\1", result_raw, flags=re.DOTALL
+                ).strip()
+                result = json.loads(json_str)
+            except json.JSONDecodeError:
+                raise e
+        return result
+
+    def _build_request(
+        self,
+        prompt: str,
+        photos: list[Photo],
+        model_id: str,
+    ) -> LlmRequest:
+        """Build a valid request that complies with openai API"""
+        return LlmRequest(
+            model=model_id,
+            temperature=self.temperature,
+            messages=[
+                # System prompt
+                LlmRequestMessage(
+                    role=LlmMessageRole.system,
+                    content=self.system_prompt,
+                ),
+                # User prompt
+                LlmRequestMessage(
+                    role=LlmMessageRole.user,
+                    content=[
+                        # Text prompt
+                        LlmMessageContent(
+                            type=LlmMessageType.text,
+                            text=prompt,
+                        ),
+                        # Image contents
+                        *[
+                            LlmMessageContent(
+                                type=LlmMessageType.image_url,
+                                image_url=photo.to_llm_image_url(),
+                            )
+                            for photo in photos
+                        ],
+                    ],
+                ),
+            ],
+        )

src/photools3/models.py

@@ -1,23 +1,43 @@
-from typing import Literal
+from enum import StrEnum
+from typing import Annotated, Any
 
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, BeforeValidator, model_validator
 
 
 class LlmRequest(BaseModel):
-    model_id: int
+    model: str = "gpt-4o"
     messages: list[LlmRequestMessage]
     temperature: float | None = None
 
 
+class LlmMessageRole(StrEnum):
+    system = "system"
+    user = "user"
+
+
 class LlmRequestMessage(BaseModel):
-    role: Literal["system", "user"]
+    role: LlmMessageRole
-    content: list[LlmMessageContent]
+    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: Literal["text", "image_url"]
+    type: LlmMessageType
     text: str | None = None
-    image_url: str | None = None
+    image_url: Base64ImageUrl | None = None
 
     @model_validator(mode="after")
     def _type_matches_content(self):

src/photools3/photo.py

@@ -4,6 +4,8 @@ from pathlib import Path
 
 from PIL import Image
 
+from photools3.models import transform_to_image_dict
+
 
 class Photo:
 
@@ -26,3 +28,6 @@ class Photo:
         self.image.save(buffer, format="JPEG", quality=85)
         image_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
         return image_data
+
+    def to_llm_image_url(self):
+        return transform_to_image_dict(self.resize().to_rgb().get_encoded())

test/conftest.py

@@ -0,0 +1,20 @@
+import pytest
+
+from photools3.llm_session import LlmSession
+from photools3.photo import Photo
+from test.test_photo import PATHS_DNG_EXAMPLES
+
+
+@pytest.fixture
+def test_photo():
+    return Photo(PATHS_DNG_EXAMPLES[0])
+
+
+@pytest.fixture
+def llm_session():
+    return LlmSession(
+        system_prompt=(
+            "You are a helpful assistent to photographer. "
+            'Answer in JSON format: `{"caption": "...", "tags": ["tag1", "tag2", ...]}`'
+        )
+    )

test/test_llm_session_integration.py

@@ -0,0 +1,13 @@
+import pytest
+
+from photools3.llm_session import LlmSession
+
+
+@pytest.mark.skipif(
+    not LlmSession(system_prompt="").is_api_reachable(), reason="API not available"
+)
+def test_llm_session(llm_session, test_photo):
+    result = llm_session.call(
+        prompt="keywords for the attached picture", photos=[test_photo]
+    )
+    assert result

test/test_llm_session_mocked.py

@@ -0,0 +1,47 @@
+from unittest.mock import MagicMock, patch
+
+from photools3.llm_session import LlmSession
+
+
+def test_llm_session_mocked(test_photo):
+    # We want to mock OpenAI client's chat.completions.create
+    with patch("photools3.llm_session.OpenAI") as mocked_openai:
+        # Setup the mock response structure
+        mock_client = MagicMock()
+        mocked_openai.return_value = mock_client
+
+        # Mock the models list to return a proper object with an ID string
+        # This prevents Pydantic ValidationError when LlmSession tries to use model_id
+        mock_model = MagicMock()
+        mock_model.id = "gpt-4o"
+        mock_client.models.list.return_value.data = [mock_model]
+
+        # Create a response object: response.choices[0].message.content
+        mock_response = MagicMock()
+        mock_choice = MagicMock()
+        mock_message = MagicMock()
+
+        mock_message.content = (
+            '{"caption": "A beautiful sunset", "tags": ["sunset", "nature"]}'
+        )
+        mock_choice.message = mock_message
+        mock_response.choices = [mock_choice]
+
+        # Configure the client to return our mock response
+        mock_client.chat.completions.create.return_value = mock_response
+
+        # Create a new LlmSession instance that will use the mocked OpenAI client
+        mocked_session = LlmSession(system_prompt="You are a helper.")
+
+        # Run the test
+        result = mocked_session.call(
+            prompt="keywords for the attached picture", photos=[test_photo]
+        )
+
+        # Verify the results
+        assert result["caption"] == "A beautiful sunset"
+        assert "sunset" in result["tags"]
+        assert "nature" in result["tags"]
+
+        # Verify that the mock was called
+        assert mock_client.chat.completions.create.called