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
-from openai.types.chat import ChatCompletion
+import json
+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:
-        return self.client.chat.completions.create(**request.model_dump())
+    def call(
+        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"]
-    content: list[LlmMessageContent]
+    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: 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())