From c5e0f987d36225b8e8b2a9355a55945fe880b1a8 Mon Sep 17 00:00:00 2001
From: jekalmin <min.jekal@navercorp.com>
Date: Mon, 25 Dec 2023 01:11:20 +0900
Subject: [PATCH] [#43] add "query_image" service

---
 .../extended_openai_conversation/__init__.py  | 65 ++++++++++++++++++-
 .../services.yaml                             | 30 +++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)
 create mode 100644 custom_components/extended_openai_conversation/services.yaml

diff --git a/custom_components/extended_openai_conversation/__init__.py b/custom_components/extended_openai_conversation/__init__.py
index 68c4d39..f0e93a3 100644
--- a/custom_components/extended_openai_conversation/__init__.py
+++ b/custom_components/extended_openai_conversation/__init__.py
@@ -5,6 +5,7 @@
 from typing import Literal
 import json
 import yaml
+import voluptuous as vol
 
 import openai
 from openai import error
@@ -12,7 +13,13 @@
 from homeassistant.components import conversation
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, MATCH_ALL, ATTR_NAME
-from homeassistant.core import HomeAssistant, ServiceCall, SupportsResponse
+from homeassistant.core import (
+    HomeAssistant,
+    ServiceCall,
+    ServiceResponse,
+    SupportsResponse,
+)
+from homeassistant.helpers.typing import ConfigType
 from homeassistant.util import ulid
 from homeassistant.components.homeassistant.exposed_entities import async_should_expose
 from homeassistant.exceptions import (
@@ -27,6 +34,7 @@
     intent,
     template,
     entity_registry as er,
+    selector,
 )
 
 from .const import (
@@ -90,6 +98,61 @@
 
 # hass.data key for agent.
 DATA_AGENT = "agent"
+SERVICE_QUERY_IMAGE = "query_image"
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up OpenAI Conversation."""
+
+    async def query_image(call: ServiceCall) -> ServiceResponse:
+        """Query an image."""
+        try:
+            model = call.data["model"]
+            images = [
+                {"type": "image_url", "image_url": image}
+                for image in call.data["images"]
+            ]
+
+            messages = [
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": call.data["prompt"]}] + images,
+                }
+            ]
+            _LOGGER.info("Prompt for %s: %s", model, messages)
+
+            response = await openai.ChatCompletion.acreate(
+                api_key=hass.data[DOMAIN][call.data["config_entry"]]["api_key"],
+                model=model,
+                messages=messages,
+                max_tokens=call.data["max_tokens"],
+            )
+            _LOGGER.info("Response %s", response)
+        except error.OpenAIError as err:
+            raise HomeAssistantError(f"Error generating image: {err}") from err
+
+        return response
+
+    hass.services.async_register(
+        DOMAIN,
+        SERVICE_QUERY_IMAGE,
+        query_image,
+        schema=vol.Schema(
+            {
+                vol.Required("config_entry"): selector.ConfigEntrySelector(
+                    {
+                        "integration": DOMAIN,
+                    }
+                ),
+                vol.Required("model", default="gpt-4-vision-preview"): cv.string,
+                vol.Required("prompt"): cv.string,
+                vol.Required("images"): vol.All(cv.ensure_list, [{"url": cv.url}]),
+                vol.Optional("max_tokens", default=300): cv.positive_int,
+            }
+        ),
+        supports_response=SupportsResponse.ONLY,
+    )
+    return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
diff --git a/custom_components/extended_openai_conversation/services.yaml b/custom_components/extended_openai_conversation/services.yaml
new file mode 100644
index 0000000..a710f23
--- /dev/null
+++ b/custom_components/extended_openai_conversation/services.yaml
@@ -0,0 +1,30 @@
+query_image:
+  fields:
+    config_entry:
+      required: true
+      selector:
+        config_entry:
+          integration: extended_openai_conversation
+    model:
+      example: gpt-4-vision-preview
+      selector:
+        text:
+    prompt:
+      example: "What’s in this image?"
+      required: true
+      selector:
+        text:
+          multiline: true
+    images:
+      example: '{"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"}'
+      required: true
+      default: []
+      selector:
+        object:
+    max_tokens:
+      example: 300
+      default: 300
+      selector:
+        number:
+          min: 1
+          mode: box