diff --git a/src/photools3/llm_session.py b/src/photools3/llm_session.py index 0ca5d3f..b850401 100644 --- a/src/photools3/llm_session.py +++ b/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 diff --git a/src/photools3/models.py b/src/photools3/models.py index 0f35a81..260c114 100644 --- a/src/photools3/models.py +++ b/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): diff --git a/src/photools3/photo.py b/src/photools3/photo.py index 1fe3874..ed5cf32 100644 --- a/src/photools3/photo.py +++ b/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()) diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/test_llm_session.py b/test/test_llm_session.py new file mode 100644 index 0000000..eb1e4a0 --- /dev/null +++ b/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