Add chat API changes

.github/workflows/ci.yml

@@ -15,5 +15,5 @@ jobs:
     uses: ballerina-platform/ballerina-library/.github/workflows/build-connector-template.yml@main
     secrets: inherit
     with:
-      repo-name: module-ballerinax-ai.model.provider.ollama
+      repo-name: module-ballerinax-ai.ollama
       publish-required: true

.github/workflows/daily-build.yml

@@ -11,4 +11,4 @@ jobs:
     uses: ballerina-platform/ballerina-library/.github/workflows/daily-build-connector-template.yml@main
     secrets: inherit
     with:
-      repo-name: module-ballerinax-ai.model.provider.ollama
+      repo-name: module-ballerinax-ai.ollama

.github/workflows/release.yml

@@ -12,5 +12,5 @@ jobs:
     uses: ballerina-platform/ballerina-library/.github/workflows/release-package-connector-template.yml@main
     secrets: inherit
     with:
-      package-name: ai.model.provider.ollama
+      package-name: ai.ollama
       package-org: ballerinax

README.md

@@ -1,7 +1,7 @@
 # Ballerina Ollama Model Provider Library
 
-[![Build](https://github.com/ballerina-platform/module-ballerinax-ai.model.provider.ollama/workflows/CI/badge.svg)](https://github.com/ballerina-platform/module-ballerinax-ai.model.provider.ollama/actions?query=workflow%3ACI)
-[![GitHub Last Commit](https://img.shields.io/github/last-commit/ballerina-platform/module-ballerinax-ai.model.provider.ollama.svg)](https://github.com/ballerina-platform/module-ballerinax-ai.model.provider.ollama/commits/master)
+[![Build](https://github.com/ballerina-platform/module-ballerinax-ai.ollama/workflows/CI/badge.svg)](https://github.com/ballerina-platform/module-ballerinax-ai.ollama/actions?query=workflow%3ACI)
+[![GitHub Last Commit](https://img.shields.io/github/last-commit/ballerina-platform/module-ballerinax-ai.ollama.svg)](https://github.com/ballerina-platform/module-ballerinax-ai.ollama/commits/master)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
 ## Overview

ballerina/Ballerina.toml

@@ -3,10 +3,13 @@ authors = ["Ballerina"]
 distribution = "2201.12.0"
 keywords = ["AI", "Agent", "Ollama", "Model", "Provider"]
 license = ["Apache-2.0"]
-name = "ai.model.provider.ollama"
+name = "ai.ollama"
 org = "ballerinax"
-repository = "https://github.com/ballerina-platform/module-ballerinax-ai.model.provider.ollama"
+repository = "https://github.com/ballerina-platform/module-ballerinax-ai.ollama"
 version = "1.0.0"
 
-[platform.java21]
-graalvmCompatible = true
+[[platform.java21.dependency]]
+groupId = "io.ballerina.lib"
+artifactId = "ai.ollama-native"
+version = "1.0.0"
+path = "../native/build/libs/ai.ollama-native-1.0.0-SNAPSHOT.jar"

ballerina/CompilerPlugin.toml

@@ -0,0 +1,9 @@
+[plugin]
+id = "ai-ollama-compiler-plugin"
+class = "io.ballerina.lib.ai.ollama.AiOllamaCompilerPlugin"
+
+[[dependency]]
+path = "../compiler-plugin/build/libs/ai.ollama-compiler-plugin-1.0.0-SNAPSHOT.jar"
+
+[[dependency]]
+path = "../compiler-plugin/build/libs/ballerina-to-openapi-2.3.0.jar"

ballerina/README.md

@@ -8,14 +8,14 @@ Ensure that your Ollama server is running locally before using this module in yo
 
 ## Quickstart
 
-To use the `ai.model.provider.ollama` module in your Ballerina application, update the `.bal` file as follows:
+To use the `ai.ollama` module in your Ballerina application, update the `.bal` file as follows:
 
 ### Step 1: Import the module
 
-Import the `ai.model.provider.ollama;` module.
+Import the `ai.ollama` module.
 
 ```ballerina
-import ballerinax/ai.model.provider.ollama;
+import ballerinax/ai.ollama;
 ```
 
 ### Step 2: Intialize the Model Provider
@@ -24,14 +24,14 @@ Here's how to initialize the Model Provider:
 
 ```ballerina
 import ballerina/ai;
-import ballerinax/ai.model.provider.ollama;
+import ballerinax/ai.ollama;
 
-final ai:ModelProvider ollamaModel = check new ollama:Provider("ollamaModelName");
+final ai:ModelProvider ollamaModel = check new ollama:ModelProvider("ollamaModelName");
 ```
 
 ### Step 4: Invoke chat completion
 
-```
+```ballerina
 ai:ChatMessage[] chatMessages = [{role: "user", content: "hi"}];
 ai:ChatAssistantMessage response = check ollamaModel->chat(chatMessages, tools = []);
 

ballerina/build.gradle

@@ -22,13 +22,15 @@ plugins{
     id 'io.ballerina.plugin'
 }
 
-description = 'Azure Model Provider - Ballerina'
+description = 'Ollama Model Provider - Ballerina'
 
-def packageName = "ai.model.provider.ollama"
+def packageName = "ai.ollama"
 def packageOrg = "ballerinax"
 def tomlVersion = stripBallerinaExtensionVersion("${project.version}")
 def ballerinaTomlFilePlaceHolder = new File("${project.rootDir}/build-config/resources/Ballerina.toml")
+def compilerPluginTomlFilePlaceHolder = new File("${project.rootDir}/build-config/resources/CompilerPlugin.toml")
 def ballerinaTomlFile = new File("$project.projectDir/Ballerina.toml")
+def compilerPluginTomlFile = new File("$project.projectDir/CompilerPlugin.toml")
 
 def stripBallerinaExtensionVersion(String extVersion) {
     if (extVersion.matches(project.ext.timestampedVersionRegex)) {
@@ -57,6 +59,11 @@ task updateTomlFiles {
         def newBallerinaToml = ballerinaTomlFilePlaceHolder.text.replace("@project.version@", project.version)
         newBallerinaToml = newBallerinaToml.replace("@toml.version@", tomlVersion)
         ballerinaTomlFile.text = newBallerinaToml
+
+        def ballerinaToOpenApiVersion = project.ballerinaToOpenApiVersion
+        def newCompilerPluginToml = compilerPluginTomlFilePlaceHolder.text.replace("@project.version@", project.version)
+        newCompilerPluginToml = newCompilerPluginToml.replace("@ballerinaToOpenApiVersion.version@", ballerinaToOpenApiVersion)
+        compilerPluginTomlFile.text = newCompilerPluginToml
     }
 }
 
@@ -95,6 +102,8 @@ clean {
     delete 'build'
 }
 
+build.dependsOn ':ai.ollama-native:build'
 build.dependsOn "generatePomFileForMavenPublication"
+build.dependsOn ":ai.ollama-compiler-plugin:build"
 publishToMavenLocal.dependsOn build
 publish.dependsOn build

ballerina/provider.bal

@@ -17,12 +17,13 @@
 import ballerina/ai;
 import ballerina/data.jsondata;
 import ballerina/http;
+import ballerina/jballerina.java;
 
 const DEFAULT_OLLAMA_SERVICE_URL = "http://localhost:11434";
 const TOOL_ROLE = "tool";
 
 # Provider represents a client for interacting with an Ollama language models.
-public isolated client class Provider {
+public isolated client class ModelProvider {
     *ai:ModelProvider;
     private final http:Client ollamaClient;
     private final string modelType;
@@ -42,7 +43,7 @@
         http:ClientConfiguration clientConfig = {...connectionConfig};
         http:Client|error ollamaClient = new (serviceUrl, clientConfig);
         if ollamaClient is error {
-            return error ai:Error("Error while connecting to the model", ollamaClient);
+            return error("Error while connecting to the model", ollamaClient);
         }
         self.modleParameters = check getModelParameterMap(modleParameters);
         self.ollamaClient = ollamaClient;
@@ -51,47 +52,82 @@
 
     # Sends a chat request to the Ollama model with the given messages and tools.
     #
-    # + messages - List of chat messages 
+    # + messages - List of chat messages or user message
     # + tools - Tool definitions to be used for the tool call
     # + stop - Stop sequence to stop the completion
     # + return - Function to be called, chat response or an error in-case of failures
-    isolated remote function chat(ai:ChatMessage[] messages, ai:ChatCompletionFunctions[] tools = [], string? stop = ())
-        returns ai:ChatAssistantMessage|ai:LlmError {
+    isolated remote function chat(ai:ChatMessage[]|ai:ChatUserMessage messages, ai:ChatCompletionFunctions[] tools = [],
+            string? stop = ()) returns ai:ChatAssistantMessage|ai:Error {
         // Ollama chat completion API reference: https://github.com/ollama/ollama/blob/main/docs/api.md#generate-a-chat-completion
-        json requestPayload = self.prepareRequestPayload(messages, tools, stop);
+        json requestPayload = check self.prepareRequestPayload(messages, tools, stop);
         OllamaResponse|error response = self.ollamaClient->/api/chat.post(requestPayload);
         if response is error {
-            return error ai:LlmConnectionError("Error while connecting to ollama", response);
+            return error("Error while connecting to ollama", response);
         }
         return self.mapOllamaResponseToAssistantMessage(response);
     }
 
-    private isolated function prepareRequestPayload(ai:ChatMessage[] messages, ai:ChatCompletionFunctions[] tools,
-            string? stop) returns json {
-        json[] transformedMessages = messages.'map(isolated function(ai:ChatMessage message) returns json {
-            if message is ai:ChatFunctionMessage {
-                return {role: TOOL_ROLE, content: message?.content};
-            }
-            return message;
-        });
+    # Sends a chat request to the model and generates a value that belongs to the type
+    # corresponding to the type descriptor argument.
+    # 
+    # + prompt - The prompt to use in the chat messages
+    # + td - Type descriptor specifying the expected return type format
+    # + return - Generates a value that belongs to the type, or an error if generation fails
+    isolated remote function generate(ai:Prompt prompt, typedesc<anydata> td = <>) returns td|ai:Error = @java:Method {
+        'class: "io.ballerina.lib.ai.ollama.Generator"
+    } external;
 
+    private isolated function prepareRequestPayload(ai:ChatMessage[]|ai:ChatUserMessage messages,
+            ai:ChatCompletionFunctions[] tools, string? stop) returns json|ai:Error {
         map<json> options = {...self.modleParameters};
         if stop is string {
             options["stop"] = [stop];
         }
 
         map<json> payload = {
             model: self.modelType,
-            messages: transformedMessages,
+            messages: check self.mapToOllamaRequestMessage(messages),
             'stream: false,
             options
         };
         if tools.length() > 0 {
-            payload["tools"] = tools.'map(tool => {'type: ai:FUNCTION, 'function: tool});
+            payload["tools"] = tools.'map(tool => {'type: FUNCTION, 'function: tool});
         }
         return payload;
     }
 
+    private isolated function mapToOllamaRequestMessage(ai:ChatMessage[]|ai:ChatUserMessage messages)
+    returns json[]|ai:Error {
+        json[] transformedMessages = [];
+        if messages is ai:ChatUserMessage {
+            transformedMessages.push({
+                role: TOOL_ROLE,
+                content: check getChatMessageStringContent(messages?.content)
+            });
+            return transformedMessages;
+        }
+        foreach ai:ChatMessage message in messages {
+            if message is ai:ChatFunctionMessage {
+                transformedMessages.push({role: TOOL_ROLE, content: message?.content});
+
+            } else if message is ai:ChatUserMessage {
+                transformedMessages.push({
+                    role: ai:USER,
+                    content: check getChatMessageStringContent(message.content)
+                });
+
+            } else if message is ai:ChatSystemMessage {
+                transformedMessages.push({
+                    role: ai:SYSTEM,
+                    content: check getChatMessageStringContent(message.content)
+                });
+            } else if message is ai:ChatAssistantMessage {
+                transformedMessages.push(message);
+            }
+        }
+        return transformedMessages;
+    }
+
     private isolated function mapOllamaResponseToAssistantMessage(OllamaResponse response)
         returns ai:ChatAssistantMessage {
         OllamaToolCall[]? toolCalls = response.message?.tool_calls;
@@ -118,6 +154,47 @@
         map<json> & readonly readonlyOptions = check options.cloneWithType();
         return readonlyOptions;
     } on fail error e {
-        return error ai:Error("Error while processing model parameters", e);
+        return error("Error while processing model parameters", e);
+    }
+}
+
+isolated function getChatMessageStringContent(ai:Prompt|string prompt) returns string|ai:Error {
+    if prompt is string {
+        return prompt;
+    }
+    string[] & readonly strings = prompt.strings;
+    anydata[] insertions = prompt.insertions;
+    string promptStr = strings[0];
+    foreach int i in 0 ..< insertions.length() {
+        string str = strings[i + 1];
+        anydata insertion = insertions[i];
+
+        if insertion is ai:TextDocument|ai:TextChunk {
+            promptStr += insertion.content + " " + str;
+            continue;
+        }
+
+        if insertion is ai:TextDocument[] {
+            foreach ai:TextDocument doc in insertion {
+                promptStr += doc.content + " ";
+            }
+            promptStr += str;
+            continue;
+        }
+
+        if insertion is ai:TextChunk[] {
+            foreach ai:TextChunk doc in insertion {
+                promptStr += doc.content + " ";
+            }
+            promptStr += str;
+            continue;
+        }
+
+        if insertion is ai:Document {
+            return error ai:Error("Only Text Documents are currently supported.");
+        }
+
+        promptStr += insertion.toString() + str;
     }
+    return promptStr.trim();
 }