src/photools3/llm_session.py

@@ -1,7 +1,17 @@
+import json
+import re
+
 from openai import OpenAI
 from openai.types.chat import ChatCompletion
 
-from photools3.models import LlmRequest
+from photools3.models import (
+    LlmMessageContent,
+    LlmMessageRole,
+    LlmMessageType,
+    LlmRequest,
+    LlmRequestMessage,
+)
+from photools3.photo import Photo
 from photools3.settings import Settings
 
 settings = Settings()
@@ -11,10 +21,11 @@ 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,
@@ -27,5 +38,55 @@ class LlmSession:
             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,
+    ) -> ChatCompletion:
+        if not model_id:
+            model_id = self.models[0].id
+
+        request = 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
+                        ],
+                    ],
+                ),
+            ],
+        )
+        response = self.client.chat.completions.create(**request.model_dump())
+        result_raw = response.choices[0].message.content
+        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

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())

test/test_llm_session.py

@@ -0,0 +1,14 @@
+from photools3.llm_session import LlmSession
+from photools3.photo import Photo
+from test.test_photo import PATHS_DNG_EXAMPLES
+
+llm_session = LlmSession(system_prompt="You are a helpful assistent to photographer")
+
+test_photo = Photo(PATHS_DNG_EXAMPLES[0])
+
+
+def test_llm_session():
+    result = llm_session.call(
+        prompt="keywords for the attaches picture", photos=[test_photo]
+    )
+    assert result