Skip to content

Commit 3e9c931

Browse files
committed
Vertex - fixing finish reason and usage info. Closes #95, #96
1 parent deb1ed9 commit 3e9c931

File tree

1 file changed

+10
-9
lines changed
  • google-vertexai-client/src/main/scala/io/cequence/openaiscala/vertexai/service/impl

1 file changed

+10
-9
lines changed

google-vertexai-client/src/main/scala/io/cequence/openaiscala/vertexai/service/impl/package.scala

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import java.{util => ju}
3838
import scala.collection.convert.ImplicitConversions.`iterable asJava`
3939
import scala.collection.convert.ImplicitConversions.`map AsJavaMap`
4040
import scala.collection.convert.ImplicitConversions.`list asScalaBuffer`
41+
import io.cequence.openaiscala.domain.response.PromptTokensDetails
4142

4243
package object impl {
4344

@@ -100,7 +101,7 @@ package object impl {
100101
.addParts(0, Part.newBuilder().setText(content).build())
101102
.build()
102103

103-
case _ => throw new OpenAIScalaClientException("Unsupported message type")
104+
case x: BaseMessage => throw new OpenAIScalaClientException(s"Unsupported message type: ${x.getClass().getName()}")
104105
}
105106

106107
def toSystemVertexAI(
@@ -267,7 +268,7 @@ package object impl {
267268
index = candidate.getIndex,
268269
message = toOpenAIAssistantMessage(candidate.getContent),
269270
finish_reason =
270-
Some(candidate.getFinishReason.getNumber + ": " + candidate.getFinishMessage),
271+
Some(candidate.getFinishReason.toString()),
271272
logprobs = None
272273
)
273274
},
@@ -288,13 +289,13 @@ package object impl {
288289
OpenAIUsageInfo(
289290
prompt_tokens = usageInfo.getPromptTokenCount,
290291
total_tokens = usageInfo.getTotalTokenCount,
291-
completion_tokens = Some(usageInfo.getPromptTokenCount)
292-
// prompt_tokens_details = Some(
293-
// PromptTokensDetails(
294-
// cached_tokens = usageInfo.getCachedContentTokenCount.getOrElse(0), TODO: add once available
295-
// audio_tokens = 0
296-
// )
297-
// )
292+
completion_tokens = Some(usageInfo.getCandidatesTokenCount),
293+
// prompt_tokens_details = Some(
294+
// PromptTokensDetails(
295+
// cached_tokens = usageInfo.getCachedContentTokenCount.getOrElse(0),
296+
// audio_tokens = None
297+
// )
298+
// )
298299
)
299300
}
300301
}

0 commit comments

Comments
 (0)