Repository: deepjavalibrary/djl
Pull request #3731 (2025-07-11 00:08:59)

Improve function call for genai extension
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/extensions/genai/src/main/java/ai/djl/genai/FunctionUtils.java b/extensions/genai/src/main/java/ai/djl/genai/FunctionUtils.java
index 1df55c919..11a567a60 100644
--- a/extensions/genai/src/main/java/ai/djl/genai/FunctionUtils.java
+++ b/extensions/genai/src/main/java/ai/djl/genai/FunctionUtils.java
@@ -13,14 +13,15 @@
 package ai.djl.genai;
 
 import ai.djl.util.JsonUtils;
+import ai.djl.util.Pair;
 import ai.djl.util.PairList;
 
-import com.google.gson.JsonElement;
-import com.google.gson.JsonObject;
+import com.google.gson.reflect.TypeToken;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
+import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -30,6 +31,7 @@ import java.util.concurrent.ConcurrentHashMap;
 public final class FunctionUtils {
 
     private static final Map<String, String> TYPE_MAPPING = new ConcurrentHashMap<>();
+    private static final Type GENERIC_MAP = new TypeToken<Map<String, String>>() {}.getType();
 
     static {
         TYPE_MAPPING.put("java.lang.String", "string");
@@ -47,6 +49,23 @@ public final class FunctionUtils {
 
     private FunctionUtils() {}
 
+    /**
+     * Returns the LLM model's function call data type.
+     *
+     * @param type the java type
+     * @return the LLM model's function call data type
+     */
+    public static String toModelType(Class<?> type) {
+        if (type.isArray()) {
+            return "array";
+        }
+        String modelType = TYPE_MAPPING.get(type.getName());
+        if (modelType != null) {
+            return modelType;
+        }
+        return "object";
+    }
+
     /**
      * Invokes the underlying method represented by the {@code Method} object.
      *
@@ -60,25 +79,33 @@ public final class FunctionUtils {
      */
     public static Object invoke(Method method, Object obj, String arguments)
             throws InvocationTargetException, IllegalAccessException {
-        Map<String, String> types = FunctionUtils.getParameters(method, false).toMap();
-        JsonObject args = JsonUtils.GSON.fromJson(arguments, JsonObject.class);
+        Map<String, Object> args = JsonUtils.GSON.fromJson(arguments, GENERIC_MAP);
+        return invoke(method, obj, args);
+    }
+
+    /**
+     * Invokes the underlying method represented by the {@code Method} object.
+     *
+     * @param method the method object
+     * @param obj the object the underlying method is invoked from
+     * @param args the arguments used for the method call
+     * @return the object returned by the method
+     * @throws IllegalAccessException if this {@code Method} object is enforcing Java language
+     *     access control and the underlying method is inaccessible.
+     * @throws IllegalArgumentException if the method is not accessible
+     */
+    public static Object invoke(Method method, Object obj, Map<String, Object> args)
+            throws InvocationTargetException, IllegalAccessException {
+        PairList<String, String> types = FunctionUtils.getParameters(method, false);
         List<Object> values = new ArrayList<>();
-        for (Map.Entry<String, JsonElement> entry : args.entrySet()) {
-            String type = types.get(entry.getKey());
-            String value = entry.getValue().getAsString();
-            if ("java.lang.String".equals(type)) {
-                values.add(value);
-            } else if ("int".equals(type) || "java.lang.Integer".equals(type)) {
-                values.add(Integer.valueOf(value));
-            } else if ("long".equals(type) || "java.lang.Long".equals(type)) {
-                values.add(Long.valueOf(value));
-            } else if ("float".equals(type) || "java.lang.Float".equals(type)) {
-                values.add(Float.valueOf(value));
-            } else if ("double".equals(type) || "java.lang.Double".equals(type)) {
-                values.add(Double.valueOf(value));
-            } else {
-                throw new IllegalArgumentException("Unsupported parameter type " + type);
+        for (Pair<String, String> pair : types) {
+            String name = pair.getKey();
+            String type = pair.getValue();
+            Object value = args.get(name);
+            if (value == null) {
+                throw new IllegalArgumentException("Missing argument: " + name);
             }
+            addArgument(values, type, value.toString());
         }
         return method.invoke(obj, values.toArray());
     }
@@ -104,19 +131,28 @@ public final class FunctionUtils {
                                 + " manually.");
             }
             String parameterName = parameter.getName();
-            String typeName = parameter.getType().getName();
-            String type = mapping ? TYPE_MAPPING.get(typeName) : typeName;
-            if (type == null) {
-                throw new IllegalArgumentException(
-                        "Unsupported parameter type "
-                                + typeName
-                                + " for parameter "
-                                + parameterName
-                                + ". Currently, only String and Java primitive types are"
-                                + " supported");
-            }
-            list.add(parameterName, type);
+            Class<?> type = parameter.getType();
+            String typeName = mapping ? toModelType(type) : type.getName();
+            list.add(parameterName, typeName);
         }
         return list;
     }
+
+    private static void addArgument(List<Object> values, String type, String value) {
+        if ("java.lang.String".equals(type)) {
+            values.add(value);
+        } else if ("boolean".equals(type) || "java.lang.Boolean".equals(type)) {
+            values.add(Boolean.valueOf(value));
+        } else if ("int".equals(type) || "java.lang.Integer".equals(type)) {
+            values.add(Integer.valueOf(value));
+        } else if ("long".equals(type) || "java.lang.Long".equals(type)) {
+            values.add(Long.valueOf(value));
+        } else if ("float".equals(type) || "java.lang.Float".equals(type)) {
+            values.add(Float.valueOf(value));
+        } else if ("double".equals(type) || "java.lang.Double".equals(type)) {
+            values.add(Double.valueOf(value));
+        } else {
+            throw new IllegalArgumentException("Unsupported parameter type " + type);
+        }
+    }
 }
diff --git a/extensions/genai/src/main/java/ai/djl/genai/gemini/GeminiOutput.java b/extensions/genai/src/main/java/ai/djl/genai/gemini/GeminiOutput.java
index 9329601b2..9dea7ad58 100644
--- a/extensions/genai/src/main/java/ai/djl/genai/gemini/GeminiOutput.java
+++ b/extensions/genai/src/main/java/ai/djl/genai/gemini/GeminiOutput.java
@@ -13,10 +13,11 @@
 package ai.djl.genai.gemini;
 
 import ai.djl.genai.gemini.types.Candidate;
-import ai.djl.genai.gemini.types.Content;
+import ai.djl.genai.gemini.types.FunctionCall;
 import ai.djl.genai.gemini.types.Part;
 import ai.djl.genai.gemini.types.UsageMetadata;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -57,6 +58,23 @@ public class GeminiOutput {
         return candidates.get(0);
     }
 
+    /**
+     * Returns the content parts.
+     *
+     * @return the content parts
+     */
+    public List<Part> getParts() {
+        Candidate candidate = getCandidate();
+        if (candidate == null) {
+            return Collections.emptyList();
+        }
+        List<Part> parts = candidate.getContent().getParts();
+        if (parts == null) {
+            return Collections.emptyList();
+        }
+        return parts;
+    }
+
     /**
      * Returns the usage metadata.
      *
@@ -81,21 +99,41 @@ public class GeminiOutput {
      * @return the aggregated text output
      */
     public String getTextOutput() {
-        if (candidates == null || candidates.isEmpty()) {
-            return "";
-        }
-        Content content = candidates.get(0).getContent();
-        List<Part> parts = content.getParts();
-        if (parts == null || parts.isEmpty()) {
-            return "";
-        }
-
         StringBuilder sb = new StringBuilder();
-        for (Part part : content.getParts()) {
+        for (Part part : getParts()) {
             if (part.getText() != null) {
                 sb.append(part.getText());
             }
         }
         return sb.toString();
     }
+
+    /**
+     * Returns the {@link FunctionCall}s in the response.
+     *
+     * @return the {@link FunctionCall}s in the response
+     */
+    public List<FunctionCall> getFunctionCalls() {
+        List<FunctionCall> list = new ArrayList<>();
+        for (Part part : getParts()) {
+            FunctionCall call = part.getFunctionCall();
+            if (call != null) {
+                list.add(call);
+            }
+        }
+        return list;
+    }
+
+    /**
+     * Returns the first {@link FunctionCall} in the response.
+     *
+     * @return the first {@link FunctionCall} in the response
+     */
+    public FunctionCall getFunctionCall() {
+        List<FunctionCall> list = getFunctionCalls();
+        if (list.isEmpty()) {
+            return null;
+        }
+        return list.get(0);
+    }
 }
diff --git a/extensions/genai/src/main/java/ai/djl/genai/gemini/types/FunctionDeclaration.java b/extensions/genai/src/main/java/ai/djl/genai/gemini/types/FunctionDeclaration.java
index 3000fe72e..0e2be6a7e 100644
--- a/extensions/genai/src/main/java/ai/djl/genai/gemini/types/FunctionDeclaration.java
+++ b/extensions/genai/src/main/java/ai/djl/genai/gemini/types/FunctionDeclaration.java
@@ -12,6 +12,15 @@
  */
 package ai.djl.genai.gemini.types;
 
+import ai.djl.genai.FunctionUtils;
+import ai.djl.util.Pair;
+import ai.djl.util.PairList;
+
+import java.lang.reflect.Method;
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
 /** A data class represents Gemini schema. */
 @SuppressWarnings("MissingJavadocMethod")
 public class FunctionDeclaration {
@@ -66,6 +75,25 @@ public class FunctionDeclaration {
         return new Builder();
     }
 
+    public static Builder function(Method method) {
+        Map<String, Schema> properties = new ConcurrentHashMap<>();
+        PairList<String, String> pairs = FunctionUtils.getParameters(method);
+        for (Pair<String, String> pair : pairs) {
+            Schema schema =
+                    Schema.builder()
+                            .type(Type.valueOf(pair.getValue().toUpperCase(Locale.ROOT)))
+                            .build();
+            properties.put(pair.getKey(), schema);
+        }
+        Schema parameters =
+                Schema.builder()
+                        .type(Type.OBJECT)
+                        .required(pairs.keys())
+                        .properties(properties)
+                        .build();
+        return builder().name(method.getName()).parameters(parameters);
+    }
+
     /** Builder class for {@code FunctionDeclaration}. */
     public static final class Builder {
 


Test patch:
diff --git a/extensions/genai/src/test/java/ai/djl/genai/gemini/GeminiTest.java b/extensions/genai/src/test/java/ai/djl/genai/gemini/GeminiTest.java
index 16e89cc9a..f827c10e0 100644
--- a/extensions/genai/src/test/java/ai/djl/genai/gemini/GeminiTest.java
+++ b/extensions/genai/src/test/java/ai/djl/genai/gemini/GeminiTest.java
@@ -13,6 +13,9 @@
 package ai.djl.genai.gemini;
 
 import ai.djl.ModelException;
+import ai.djl.genai.FunctionUtils;
+import ai.djl.genai.gemini.types.FunctionCall;
+import ai.djl.genai.gemini.types.FunctionDeclaration;
 import ai.djl.genai.gemini.types.GenerationConfig;
 import ai.djl.genai.gemini.types.GoogleSearch;
 import ai.djl.genai.gemini.types.HarmBlockThreshold;
@@ -31,9 +34,11 @@ import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Map;
 
 public class GeminiTest {
 
@@ -94,6 +99,53 @@ public class GeminiTest {
         }
     }
 
+    @Test
+    public void testGenerateContentWithFunction()
+            throws ModelException, IOException, TranslateException, ReflectiveOperationException {
+        String mockResponse = loadFunctionContent();
+
+        Method method = GeminiTest.class.getMethod("getWeather", String.class, boolean.class);
+        FunctionDeclaration function =
+                FunctionDeclaration.function(method)
+                        .description(
+                                "Get the current weather in a given location, set unit to true for"
+                                        + " celsius")
+                        .build();
+        GenerationConfig config =
+                GenerationConfig.builder()
+                        .candidateCount(1)
+                        .addTool(Tool.builder().addFunctionDeclaration(function))
+                        .build();
+
+        try (TestServer server = TestServer.newInstance(mockResponse)) {
+            String endpoint = "http://localhost:" + server.getPort();
+            // endpoint = "https://generativelanguage.googleapis.com";
+            String url = endpoint + "/v1beta/models/gemini-2.5-flash:generateContent";
+            Criteria<GeminiInput, GeminiOutput> criteria =
+                    Criteria.builder()
+                            .setTypes(GeminiInput.class, GeminiOutput.class)
+                            .optModelUrls(url)
+                            .optArgument("API_KEY", "1234") // override env var
+                            .build();
+
+            try (ZooModel<GeminiInput, GeminiOutput> model = criteria.loadModel();
+                    Predictor<GeminiInput, GeminiOutput> predictor = model.newPredictor()) {
+
+                GeminiInput in =
+                        GeminiInput.text("What is the weather like in celsius in New York today?")
+                                .generationConfig(config)
+                                .build();
+
+                GeminiOutput ret = predictor.predict(in);
+                FunctionCall functionCall = ret.getFunctionCall();
+                Assert.assertNotNull(functionCall);
+                Map<String, Object> arguments = functionCall.getArgs();
+                String weather = (String) FunctionUtils.invoke(method, this, arguments);
+                Assert.assertEquals(weather, "nice");
+            }
+        }
+    }
+
     @Test
     public void testImageUnderstand() throws ModelException, IOException, TranslateException {
         String mockResponse = loadGeneratedContent();
@@ -204,4 +256,14 @@ public class GeminiTest {
                 + "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\" is a test.\"}],"
                 + "\"role\": \"model\"},\"finishReason\": \"STOP\",\"index\":0}]}\n\n";
     }
+
+    private String loadFunctionContent() {
+        return "{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{"
+                + "\"args\":{\"unit\":true,\"location\":\"New York\"},\"name\":\"getWeather\"}}],"
+                + "\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}]}";
+    }
+
+    public String getWeather(String location, boolean unit) {
+        return "nice";
+    }
 }