Add text and image document support for Ballerina OpenAI model provider

ballerina/Dependencies.toml

@@ -5,7 +5,7 @@
 
 [ballerina]
 dependencies-toml-version = "2"
-distribution-version = "2201.12.7"
+distribution-version = "2201.12.0"
 
 [[package]]
 org = "ballerina"
@@ -65,6 +65,9 @@ version = "1.7.0"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"}
 ]
+modules = [
+	{org = "ballerina", packageName = "constraint", moduleName = "constraint"}
+]
 
 [[package]]
 org = "ballerina"
@@ -186,6 +189,9 @@ dependencies = [
 	{org = "ballerina", name = "jballerina.java"},
 	{org = "ballerina", name = "lang.__internal"}
 ]
+modules = [
+	{org = "ballerina", packageName = "lang.array", moduleName = "lang.array"}
+]
 
 [[package]]
 org = "ballerina"
@@ -403,8 +409,10 @@ name = "ai.openai"
 version = "1.1.0"
 dependencies = [
 	{org = "ballerina", name = "ai"},
+	{org = "ballerina", name = "constraint"},
 	{org = "ballerina", name = "http"},
 	{org = "ballerina", name = "jballerina.java"},
+	{org = "ballerina", name = "lang.array"},
 	{org = "ballerina", name = "lang.regexp"},
 	{org = "ballerina", name = "log"},
 	{org = "ballerina", name = "test"},

ballerina/provider_utils.bal

@@ -15,13 +15,20 @@
 // under the License.
 
 import ballerina/ai;
+import ballerina/constraint;
+import ballerina/lang.array;
 import ballerinax/openai.chat;
 
 type ResponseSchema record {|
     map<json> schema;
     boolean isOriginallyJsonObject = true;
 |};
 
+type DocumentContentPart TextContentPart|ImageContentPart;
+
+type TextContentPart chat:ChatCompletionRequestMessageContentPartText;
+type ImageContentPart chat:ChatCompletionRequestMessageContentPartImage;
+
 const JSON_CONVERSION_ERROR = "FromJsonStringError";
 const CONVERSION_ERROR = "ConversionError";
 const ERROR_MESSAGE = "Error occurred while attempting to parse the response from the " +
@@ -95,34 +102,95 @@ isolated function getGetResultsTool(map<json> parameters) returns chat:ChatCompl
         }
     ];
 
-isolated function generateChatCreationContent(ai:Prompt prompt) returns string|ai:Error {
+isolated function generateChatCreationContent(ai:Prompt prompt)
+                        returns DocumentContentPart[]|ai:Error {
     string[] & readonly strings = prompt.strings;
     anydata[] insertions = prompt.insertions;
-    string promptStr = strings[0];
+    DocumentContentPart[] contentParts = [];
+    string accumulatedTextContent = "";
+
+    if strings.length() > 0 {
+        accumulatedTextContent += strings[0];
+    }
+
     foreach int i in 0 ..< insertions.length() {
-        string str = strings[i + 1];
         anydata insertion = insertions[i];
+        string str = strings[i + 1];
 
-        if insertion is ai:TextDocument {
-            promptStr += insertion.content + " " + str;
-            continue;
+        if insertion is ai:Document {
+            addTextContentPart(buildTextContentPart(accumulatedTextContent), contentParts);
+            accumulatedTextContent = "";
+            check addDocumentContentPart(insertion, contentParts);
+        } else if insertion is ai:Document[] {
+            addTextContentPart(buildTextContentPart(accumulatedTextContent), contentParts);
+            accumulatedTextContent = "";
+            foreach ai:Document doc in insertion {
+                check addDocumentContentPart(doc, contentParts);
+            }
+        } else {
+            accumulatedTextContent += insertion.toString();
         }
+        accumulatedTextContent += str;
+    }
 
-        if insertion is ai:TextDocument[] {
-            foreach ai:TextDocument doc in insertion {
-                promptStr += doc.content  + " ";
-            }
-            promptStr += str;
-            continue;
+    addTextContentPart(buildTextContentPart(accumulatedTextContent), contentParts);
+    return contentParts;
+}
+
+isolated function addDocumentContentPart(ai:Document doc, DocumentContentPart[] contentParts) returns ai:Error? {
+    if doc is ai:TextDocument {
+        return addTextContentPart(buildTextContentPart(doc.content), contentParts);
+    } else if doc is ai:ImageDocument {
+        return contentParts.push(check buildImageContentPart(doc));
+    }
+
+    return error ai:Error("Only text and image documents are supported.");
+}
+
+isolated function addTextContentPart(TextContentPart? contentPart, DocumentContentPart[] contentParts) {
+    if contentPart is TextContentPart {
+        return contentParts.push(contentPart);
+    }
+}
+
+isolated function buildTextContentPart(string content) returns TextContentPart? {
+    if content.length() == 0 {
+        return;
+    }
+
+    return {
+        'type: "text",
+        text: content
+    };
+}
+
+isolated function buildImageContentPart(ai:ImageDocument doc) returns ImageContentPart|ai:Error =>
+    {
+        'type: "image_url",
+        image_url: {
+            url: check buildImageUrl(doc.content, doc.metadata?.mimeType)
         }
+    };
 
-        if insertion is ai:Document {
-            return error ai:Error("Only Text Documents are currently supported.");
+isolated function buildImageUrl(ai:Url|byte[] content, string? mimeType) returns string|ai:Error {
+    if content is ai:Url {
+        ai:Url|constraint:Error validationRes = constraint:validate(content);
+        if validationRes is error {
+            return error(validationRes.message(), validationRes.cause());
         }
+        return content;
+    }
+
+    return string `data:${mimeType ?: "image/*"};base64,${check getBase64EncodedString(content)}`;
+}
 
-        promptStr += insertion.toString() + str;
+isolated function getBase64EncodedString(byte[] content) returns string|ai:Error {
+    string|error binaryContent = array:toBase64(content);
+    if binaryContent is error {
+        return error("Failed to convert byte array to string: " + binaryContent.message() + ", " +
+                        binaryContent.detail().toBalString());
     }
-    return promptStr.trim();
+    return binaryContent;
 }
 
 isolated function handleParseResponseError(error chatResponseError) returns error {
@@ -135,7 +203,7 @@ isolated function handleParseResponseError(error chatResponseError) returns erro
 
 isolated function generateLlmResponse(chat:Client llmClient, OPEN_AI_MODEL_NAMES modelType, 
         ai:Prompt prompt, typedesc<json> expectedResponseTypedesc) returns anydata|ai:Error {
-    string content = check generateChatCreationContent(prompt);
+    DocumentContentPart[] content = check generateChatCreationContent(prompt);
     ResponseSchema ResponseSchema = check getExpectedResponseSchema(expectedResponseTypedesc);
     chat:ChatCompletionTool[]|error tools = getGetResultsTool(ResponseSchema.schema);
     if tools is error {
@@ -157,7 +225,7 @@ isolated function generateLlmResponse(chat:Client llmClient, OPEN_AI_MODEL_NAMES
     chat:CreateChatCompletionResponse|error response =
         llmClient->/chat/completions.post(request);
     if response is error {
-        return error("LLM call failed: " + response.message());
+        return error("LLM call failed: " + response.message(), detail = response.detail(), cause = response.cause());
     }
 
     chat:CreateChatCompletionResponse_choices[] choices = response.choices;

ballerina/tests/test_services.bal

@@ -19,20 +19,25 @@ import ballerina/test;
 import ballerinax/openai.chat;
 
 service /llm on new http:Listener(8080) {
-    resource function post openai/chat/completions(chat:CreateChatCompletionRequest payload)
+    // Change the payload type to JSON due to https://github.com/ballerina-platform/ballerina-library/issues/8048.
+    resource function post openai/chat/completions(@http:Payload json payload)
                 returns chat:CreateChatCompletionResponse|error {
         test:assertEquals(payload.model, GPT_4O);
-        chat:ChatCompletionRequestMessage[] messages = check payload.messages.ensureType();
+        chat:ChatCompletionRequestMessage[] messages = check (check payload.messages).fromJsonWithType();
         chat:ChatCompletionRequestMessage message = messages[0];
 
-        string? content = check message["content"].ensureType();
+        chat:ChatCompletionRequestUserMessageContentPart[]? content = check message["content"].ensureType();
         if content is () {
             test:assertFail("Expected content in the payload");
         }
 
-        test:assertEquals(content, getExpectedPrompt(content));
+        chat:ChatCompletionRequestUserMessageContentPart initialContentPart = content[0];
+        TextContentPart initialTextContent = check initialContentPart.ensureType();
+        string initialText = initialTextContent.text;
+        test:assertEquals(content, getExpectedContentParts(initialText),
+                string `Test failed for prompt with initial content, ${initialText}`);
         test:assertEquals(message.role, "user");
-        chat:ChatCompletionTool[]? tools = payload.tools;
+        chat:ChatCompletionTool[]? tools = check (check payload.tools).fromJsonWithType();
         if tools is () || tools.length() == 0 {
             test:assertFail("No tools in the payload");
         }
@@ -42,7 +47,8 @@ service /llm on new http:Listener(8080) {
             test:assertFail("No parameters in the expected tool");
         }
 
-        test:assertEquals(parameters, getExpectedParameterSchema(content), string `Test failed for prompt:- ${content}`);
-        return getTestServiceResponse(content);
+        test:assertEquals(parameters, getExpectedParameterSchema(initialText),
+                string `Test failed for prompt with initial content, ${initialText}`);
+        return getTestServiceResponse(initialText);
     }
 }