rename "initalMessage" add token info to Response interface

lucee · Aug 29, 2024 · 8461533 · 8461533
1 parent 4b7013d
commit 8461533
Show file tree

Hide file tree

Showing 12 changed files with 78 additions and 18 deletions.
diff --git a/core/src/main/java/lucee/runtime/PageContextImpl.java b/core/src/main/java/lucee/runtime/PageContextImpl.java
@@ -4214,12 +4214,12 @@ public AIEngine getAIEngine(String nameAI, AIEngine defaultValue) {
 		}
 	}
 
-	public AISession createAISession(String nameAI, String initalMessage) throws PageException {
-		return getAIEngine(nameAI).createSession(initalMessage, -1);
+	public AISession createAISession(String nameAI, String systemMessage) throws PageException {
+		return getAIEngine(nameAI).createSession(systemMessage, -1);
 	}
 
-	public AISession createAISession(String nameAI, String initalMessage, long timeout) throws PageException {
-		return getAIEngine(nameAI).createSession(initalMessage, timeout);
+	public AISession createAISession(String nameAI, String systemMessage, long timeout) throws PageException {
+		return getAIEngine(nameAI).createSession(systemMessage, timeout);
 	}
 
 	public String getNameFromDefault(String defaultName) throws PageException {

diff --git a/core/src/main/java/lucee/runtime/ai/Response.java b/core/src/main/java/lucee/runtime/ai/Response.java
@@ -2,4 +2,6 @@
 
 public interface Response {
 	public String getAnswer();
+
+	public long getTotalTokenUsed();
 }
diff --git a/core/src/main/java/lucee/runtime/ai/google/GeminiModel.java b/core/src/main/java/lucee/runtime/ai/google/GeminiModel.java
@@ -17,4 +17,22 @@ public GeminiModel(Struct raw, String charset) throws PageException {
 	public String getLabel() {
 		return Caster.toString(raw.get(KeyConstants._displayName, null), null);
 	}
+
+	@Override
+	public String getDescription() {
+		return Caster.toString(raw.get(KeyConstants._description, null), null);
+	}
+
+	public long getInputTokenLimit() {
+		return Caster.toLongValue(raw.get("inputTokenLimit", null), 0L);
+	}
+
+	public long getOutputTokenLimit() {
+		return Caster.toLongValue(raw.get("outputTokenLimit", null), 0L);
+	}
+
+	public Struct getData() {
+		return raw;
+	}
+
 }
diff --git a/core/src/main/java/lucee/runtime/ai/google/GeminiResponse.java b/core/src/main/java/lucee/runtime/ai/google/GeminiResponse.java
@@ -15,6 +15,7 @@ public class GeminiResponse implements Response {
 
 	private Struct raw;
 	private String charset;
+	private long tokens = -1L;
 
 	public GeminiResponse(Struct raw, String charset) {
 		this.raw = raw;
@@ -55,4 +56,13 @@ public Struct getData() {
 		return raw;
 	}
 
+	@Override
+	public long getTotalTokenUsed() {
+		if (tokens == -1L) {
+			Struct sct = Caster.toStruct(raw.get("usageMetadata", null), null);
+			if (sct == null) return tokens = 0L;
+			return tokens = Caster.toLongValue(sct.get("totalTokenCount", null), 0L);
+		}
+		return tokens;
+	}
 }
diff --git a/core/src/main/java/lucee/runtime/ai/google/GeminiSession.java b/core/src/main/java/lucee/runtime/ai/google/GeminiSession.java
@@ -35,12 +35,12 @@
 
 public class GeminiSession extends AISessionSupport {
 	private GeminiEngine geminiEngine;
-	private String initalMessage;
+	private String systemMessage;
 
-	public GeminiSession(GeminiEngine engine, String initalMessage, long timeout) {
+	public GeminiSession(GeminiEngine engine, String systemMessage, long timeout) {
 		super(engine, timeout);
 		this.geminiEngine = engine;
-		this.initalMessage = initalMessage;
+		this.systemMessage = systemMessage;
 	}
 
 	@Override
@@ -59,8 +59,8 @@ public Response inquiry(String message, AIResponseListener listener) throws Page
 			root.set(KeyConstants._contents, contents);
 
 			// add system
-			if (!StringUtil.isEmpty(initalMessage, true)) {
-				contents.append(createParts("user", initalMessage));
+			if (!StringUtil.isEmpty(systemMessage, true)) {
+				contents.append(createParts("user", systemMessage));
 			}
 
 			// Add conversation history

diff --git a/core/src/main/java/lucee/runtime/ai/google/GeminiStreamResponse.java b/core/src/main/java/lucee/runtime/ai/google/GeminiStreamResponse.java
@@ -23,6 +23,8 @@ public class GeminiStreamResponse implements Response {
 
 	private AIResponseListener listener;
 
+	private long tokens = -1L;
+
 	public GeminiStreamResponse(String charset, AIResponseListener listener) {
 		this.charset = charset;
 		this.listener = listener;
@@ -68,4 +70,16 @@ public void addPart(Struct part) throws PageException {
 		answer.append(str);
 	}
 
+	@Override
+	public long getTotalTokenUsed() {
+		if (tokens == -1L) {
+			Struct sct = Caster.toStruct(raw.get(raw.size(), null), null);
+			if (sct == null) return tokens = 0L;
+
+			sct = Caster.toStruct(sct.get("usageMetadata", null), null);
+			if (sct == null) return tokens = 0L;
+			return tokens = Caster.toLongValue(sct.get("totalTokenCount", null), 0L);
+		}
+		return tokens;
+	}
 }
diff --git a/core/src/main/java/lucee/runtime/ai/openai/OpenAIResponse.java b/core/src/main/java/lucee/runtime/ai/openai/OpenAIResponse.java
@@ -15,6 +15,7 @@ public class OpenAIResponse implements Response {
 
 	private Struct raw;
 	private String charset;
+	private long tokens = -1L;
 
 	public OpenAIResponse(Struct raw, String charset) {
 		this.raw = raw;
@@ -48,4 +49,14 @@ public Struct getData() {
 		return raw;
 	}
 
+	@Override
+	public long getTotalTokenUsed() {
+		if (tokens == -1L) {
+			Struct sct = Caster.toStruct(raw.get("usage", null), null);
+			if (sct == null) return tokens = 0L;
+			return tokens = Caster.toLongValue(sct.get("total_tokens", null), 0L);
+		}
+		return tokens;
+	}
+
 }
diff --git a/core/src/main/java/lucee/runtime/ai/openai/OpenAISession.java b/core/src/main/java/lucee/runtime/ai/openai/OpenAISession.java
@@ -36,12 +36,12 @@
 public class OpenAISession extends AISessionSupport {
 
 	private OpenAIEngine openaiEngine;
-	private String initalMessage;
+	private String systemMessage;
 
-	public OpenAISession(OpenAIEngine engine, String initalMessage, long timeout) {
+	public OpenAISession(OpenAIEngine engine, String systemMessage, long timeout) {
 		super(engine, timeout);
 		this.openaiEngine = engine;
-		this.initalMessage = initalMessage;
+		this.systemMessage = systemMessage;
 	}
 
 	@Override
@@ -56,10 +56,10 @@ public Response inquiry(String message, AIResponseListener listener) throws Page
 			Array arr = new ArrayImpl();
 
 			// add system
-			if (!StringUtil.isEmpty(initalMessage)) {
+			if (!StringUtil.isEmpty(systemMessage)) {
 				msg = new StructImpl(StructImpl.TYPE_LINKED);
 				msg.set(KeyConstants._role, "system");
-				msg.set(KeyConstants._content, initalMessage);
+				msg.set(KeyConstants._content, systemMessage);
 				arr.append(msg);
 			}
 

diff --git a/core/src/main/java/lucee/runtime/ai/openai/OpenAIStreamResponse.java b/core/src/main/java/lucee/runtime/ai/openai/OpenAIStreamResponse.java
@@ -49,6 +49,7 @@ public Struct getData() {
 
 	public void addPart(Struct part) throws PageException {
 		if (raw == null) raw = part;
+
 		// raw.appendEL(part);
 		Array arr = Caster.toArray(part.get("choices", null), null);
 		// print.e(arr);
@@ -65,4 +66,8 @@ public void addPart(Struct part) throws PageException {
 		if (listener != null) listener.listen(str);
 	}
 
+	@Override
+	public long getTotalTokenUsed() {
+		return 0;
+	}
 }
diff --git a/core/src/main/java/lucee/runtime/functions/ai/CreateAISession.java b/core/src/main/java/lucee/runtime/functions/ai/CreateAISession.java
@@ -18,9 +18,9 @@ public static Object call(PageContext pc, String nameAI) throws PageException {
 		return call(pc, nameAI, null);
 	}
 
-	public static Object call(PageContext pc, String nameAI, String initalMessage) throws PageException {
+	public static Object call(PageContext pc, String nameAI, String systemMessage) throws PageException {
 		if (nameAI.startsWith("default:")) nameAI = ((PageContextImpl) pc).getNameFromDefault(nameAI.substring(8));
-		return ((PageContextImpl) pc).createAISession(nameAI, initalMessage);
+		return ((PageContextImpl) pc).createAISession(nameAI, systemMessage);
 	}
 
 	@Override

diff --git a/loader/build.xml b/loader/build.xml
@@ -2,7 +2,7 @@
 <project default="core" basedir="." name="Lucee" 
   xmlns:resolver="antlib:org.apache.maven.resolver.ant">
 
-  <property name="version" value="6.2.0.63-SNAPSHOT"/>
+  <property name="version" value="6.2.0.64-SNAPSHOT"/>
 
   <taskdef uri="antlib:org.apache.maven.resolver.ant" resource="org/apache/maven/resolver/ant/antlib.xml">
     <classpath>

diff --git a/loader/pom.xml b/loader/pom.xml
@@ -3,7 +3,7 @@
 
   <groupId>org.lucee</groupId>
   <artifactId>lucee</artifactId>
-  <version>6.2.0.63-SNAPSHOT</version>
+  <version>6.2.0.64-SNAPSHOT</version>
   <packaging>jar</packaging>
 
   <name>Lucee Loader Build</name>