LLM: Add image recognition and generation support

app/domain/__init__.py

@@ -5,3 +5,4 @@
     TutorChatPipelineExecutionDTO,
 )
 from .pyris_message import PyrisMessage, IrisMessageRole
+from .pyris_image import PyrisImage

app/domain/iris_message.py

@@ -0,0 +1,19 @@
+from enum import Enum
+from pydantic import BaseModel
+from typing import List, Optional
+from .pyris_image import PyrisImage
+
+
+class IrisMessageRole(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
+class IrisMessage(BaseModel):
+    text: str = ""
+    role: IrisMessageRole
+    images: Optional[List[PyrisImage]] = None
+
+    def __str__(self):
+        return f"{self.role.lower()}: {self.text}"

app/domain/pyris_image.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+from typing import Optional
+
+
+class PyrisImage(BaseModel):
+    base64: str
+    prompt: Optional[str] = None
+    mime_type: Optional[str] = "jpeg"
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "prompt": "Example prompt",
+                "base64": "base64EncodedString==",
+                "mime_type": "jpeg",
+            }
+        }

app/llm/external/openai_completion.py

@@ -2,6 +2,7 @@
 from openai import OpenAI
 from openai.lib.azure import AzureOpenAI
 
+from ...domain import PyrisImage
 from ...llm import CompletionArguments
 from ...llm.external.model import CompletionModel
 
@@ -11,7 +12,9 @@ class OpenAICompletionModel(CompletionModel):
     api_key: str
     _client: OpenAI
 
-    def complete(self, prompt: str, arguments: CompletionArguments) -> any:
+    def complete(
+        self, prompt: str, arguments: CompletionArguments, images: [PyrisImage] = None
+    ) -> any:
         response = self._client.completions.create(
             model=self.model,
             prompt=prompt,

app/llm/external/openai_dalle.py

@@ -0,0 +1,60 @@
+import base64
+from datetime import datetime
+from typing import Literal, Any
+
+import requests
+from openai import OpenAI
+
+from ...domain.pyris_image import PyrisImage
+from ...llm.external.model import ImageGenerationModel
+
+
+class OpenAIDalleWrapper(ImageGenerationModel):
+    type: Literal["openai_dalle"]
+    model: str
+    _client: OpenAI
+
+    def model_post_init(self, __context: Any) -> None:
+        self._client = OpenAI(api_key=self.api_key)
+
+    def generate_images(
+        self,
+        prompt: str,
+        n: int = 1,
+        size: Literal[
+            "256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"
+        ] = "256x256",
+        quality: Literal["standard", "hd"] = "standard",
+        **kwargs
+    ) -> [PyrisImage]:
+        response = self._client.images.generate(
+            model=self.model,
+            prompt=prompt,
+            size=size,
+            quality=quality,
+            n=n,
+            response_format="url",
+            **kwargs
+        )
+
+        images = response.data
+        iris_images = []
+        for image in images:
+            if image.revised_prompt is None:
+                image.revised_prompt = prompt
+            if image.b64_json is None:
+                image_response = requests.get(image.url)
+                image.b64_json = base64.b64encode(image_response.content).decode(
+                    "utf-8"
+                )
+
+            iris_images.append(
+                PyrisImage(
+                    prompt=image.revised_prompt,
+                    base64=image.b64_json,
+                    timestamp=datetime.fromtimestamp(response.created),
+                    raw_data=image,
+                )
+            )
+
+        return iris_images