[goose-llm] fix image content bug, add optional request_id field (#3439)

This commit is contained in:
Salman Mohammed
2025-07-15 18:06:37 -04:00
committed by GitHub
parent f4e3d06f9e
commit a5d77950db
22 changed files with 482 additions and 512 deletions
+115
View File
@@ -0,0 +1,115 @@
import kotlin.system.measureNanoTime
import kotlinx.coroutines.runBlocking
import uniffi.goose_llm.*
import java.net.URI
import java.net.http.HttpClient
import java.net.http.HttpRequest
import java.net.http.HttpResponse
/* ---------- Goose helpers ---------- */
fun buildProviderConfig(host: String, token: String): String =
"""{ "host": "$host", "token": "$token" }"""
suspend fun timeGooseCall(
modelCfg: ModelConfig,
providerName: String,
providerCfg: String
): Pair<Double, CompletionResponse> {
val req = createCompletionRequest(
providerName,
providerCfg,
modelCfg,
systemPreamble = "You are a helpful assistant.",
messages = listOf(
Message(
Role.USER,
System.currentTimeMillis() / 1000,
listOf(MessageContent.Text(TextContent("Write me a 1000 word chapter about learning Go vs Rust in the world of LLMs and AI.")))
)
),
extensions = emptyList()
)
lateinit var resp: CompletionResponse
val wallMs = measureNanoTime { resp = completion(req) } / 1_000_000.0
return wallMs to resp
}
/* ---------- OpenAI helpers ---------- */
fun timeOpenAiCall(client: HttpClient, apiKey: String): Double {
val body = """
{
"model": "gpt-4.1",
"max_tokens": 500,
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Write me a 1000 word chapter about learning Go vs Rust in the world of LLMs and AI."}
]
}
""".trimIndent()
val request = HttpRequest.newBuilder()
.uri(URI.create("https://api.openai.com/v1/chat/completions"))
.header("Authorization", "Bearer $apiKey")
.header("Content-Type", "application/json")
.POST(HttpRequest.BodyPublishers.ofString(body))
.build()
val wallMs = measureNanoTime {
client.send(request, HttpResponse.BodyHandlers.ofString())
} / 1_000_000.0
return wallMs
}
/* ---------- main ---------- */
fun main() = runBlocking {
/* Goose provider setup */
val providerName = "databricks"
val host = System.getenv("DATABRICKS_HOST") ?: error("DATABRICKS_HOST not set")
val token = System.getenv("DATABRICKS_TOKEN") ?: error("DATABRICKS_TOKEN not set")
val providerCfg = buildProviderConfig(host, token)
/* OpenAI setup */
val openAiKey = System.getenv("OPENAI_API_KEY") ?: error("OPENAI_API_KEY not set")
val httpClient = HttpClient.newBuilder().build()
val gooseModels = listOf("goose-claude-4-sonnet", "goose-gpt-4-1")
val runsPerModel = 3
/* --- Goose timing --- */
for (model in gooseModels) {
val maxTokens = 500
val cfg = ModelConfig(model, 100_000u, 0.0f, maxTokens)
var wallSum = 0.0
var gooseSum = 0.0
println("=== Goose: $model ===")
repeat(runsPerModel) { run ->
val (wall, resp) = timeGooseCall(cfg, providerName, providerCfg)
val gooseMs = resp.runtimeMetrics.totalTimeSec * 1_000
val overhead = wall - gooseMs
wallSum += wall
gooseSum += gooseMs
println("run ${run + 1}: wall = %.1f ms | goose-llm = %.1f ms | overhead = %.1f ms"
.format(wall, gooseMs, overhead))
}
println("-- avg wall = %.1f ms | avg overhead = %.1f ms --\n"
.format(wallSum / runsPerModel, (wallSum - gooseSum) / runsPerModel))
}
/* --- OpenAI direct timing --- */
var oaSum = 0.0
println("=== OpenAI: gpt-4.1 (direct HTTPS) ===")
repeat(runsPerModel) { run ->
val wall = timeOpenAiCall(httpClient, openAiKey)
oaSum += wall
println("run ${run + 1}: wall = %.1f ms".format(wall))
}
println("-- avg wall = %.1f ms --".format(oaSum / runsPerModel))
}
+206 -270
View File
@@ -1,292 +1,228 @@
import java.io.File
import java.util.Base64
import kotlinx.coroutines.runBlocking
import uniffi.goose_llm.*
fun main() = runBlocking {
val now = System.currentTimeMillis() / 1000
val msgs = listOf(
// 1) User sends a plain-text prompt
Message(
role = Role.USER,
created = now,
content = listOf(
MessageContent.Text(
TextContent("What is 7 x 6?")
)
)
),
/* ---------- shared helpers ---------- */
// 2) Assistant makes a tool request (ToolReq) to calculate 7×6
Message(
role = Role.ASSISTANT,
created = now + 2,
content = listOf(
MessageContent.ToolReq(
ToolRequest(
id = "calc1",
toolCall = """
{
"status": "success",
"value": {
"name": "calculator_extension__toolname",
"arguments": {
"operation": "doesnotexist",
"numbers": [7, 6]
},
"needsApproval": false
}
}
""".trimIndent()
)
)
)
),
// 3) User (on behalf of the tool) responds with the tool result (ToolResp)
Message(
role = Role.USER,
created = now + 3,
content = listOf(
MessageContent.ToolResp(
ToolResponse(
id = "calc1",
toolResult = """
{
"status": "error",
"error": "Invalid value for operation: 'doesnotexist'. Valid values are: ['add', 'subtract', 'multiply', 'divide']"
}
""".trimIndent()
)
)
)
),
// 4) Assistant makes a tool request (ToolReq) to calculate 7×6
Message(
role = Role.ASSISTANT,
created = now + 4,
content = listOf(
MessageContent.ToolReq(
ToolRequest(
id = "calc1",
toolCall = """
{
"status": "success",
"value": {
"name": "calculator_extension__toolname",
"arguments": {
"operation": "multiply",
"numbers": [7, 6]
},
"needsApproval": false
}
}
""".trimIndent()
)
)
)
),
// 5) User (on behalf of the tool) responds with the tool result (ToolResp)
Message(
role = Role.USER,
created = now + 5,
content = listOf(
MessageContent.ToolResp(
ToolResponse(
id = "calc1",
toolResult = """
{
"status": "success",
"value": [
{"type": "text", "text": "42"}
]
}
""".trimIndent()
)
)
)
),
)
printMessages(msgs)
println("---\n")
// Setup provider
val providerName = "databricks"
val host = System.getenv("DATABRICKS_HOST") ?: error("DATABRICKS_HOST not set")
val token = System.getenv("DATABRICKS_TOKEN") ?: error("DATABRICKS_TOKEN not set")
val providerConfig = """{"host": "$host", "token": "$token"}"""
println("Provider Name: $providerName")
println("Provider Config: $providerConfig")
val sessionName = generateSessionName(providerName, providerConfig, msgs)
println("\nSession Name: $sessionName")
val tooltip = generateTooltip(providerName, providerConfig, msgs)
println("\nTooltip: $tooltip")
// Completion
val modelName = "goose-gpt-4-1"
val modelConfig = ModelConfig(
modelName,
100000u, // UInt
0.1f, // Float
200 // Int
)
fun buildProviderConfig(host: String, token: String, imageFormat: String = "OpenAi"): String = """
{
"host": "$host",
"token": "$token",
"image_format": "$imageFormat"
}
""".trimIndent()
fun calculatorExtension(): ExtensionConfig {
val calculatorTool = createToolConfig(
name = "calculator",
name = "calculator",
description = "Perform basic arithmetic operations",
inputSchema = """
{
"type": "object",
"required": ["operation", "numbers"],
"properties": {
"operation": {
"type": "string",
"enum": ["add", "subtract", "multiply", "divide"],
"description": "The arithmetic operation to perform"
},
"numbers": {
"type": "array",
"items": { "type": "number" },
"description": "List of numbers to operate on in order"
}
"type": "object",
"required": ["operation", "numbers"],
"properties": {
"operation": {
"type": "string",
"enum": ["add", "subtract", "multiply", "divide"],
"description": "The arithmetic operation to perform"
},
"numbers": {
"type": "array",
"items": { "type": "number" },
"description": "List of numbers to operate on in order"
}
}
}
""".trimIndent(),
approvalMode = ToolApprovalMode.AUTO
)
val calculator_extension = ExtensionConfig(
name = "calculator_extension",
return ExtensionConfig(
name = "calculator_extension",
instructions = "This extension provides a calculator tool.",
tools = listOf(calculatorTool)
tools = listOf(calculatorTool)
)
val extensions = listOf(calculator_extension)
val systemPreamble = "You are a helpful assistant."
// Testing with tool calls with an error in tool name
val reqToolErr = createCompletionRequest(
providerName,
providerConfig,
modelConfig,
systemPreamble,
messages = listOf(
Message(
role = Role.USER,
created = now,
content = listOf(
MessageContent.Text(
TextContent("What is 7 x 6?")
)
)
)),
extensions = extensions
)
val respToolErr = completion(reqToolErr)
println("\nCompletion Response (one msg):\n${respToolErr.message}")
println()
val reqAll = createCompletionRequest(
providerName,
providerConfig,
modelConfig,
systemPreamble,
messages = msgs,
extensions = extensions
)
val respAll = completion(reqAll)
println("\nCompletion Response (all msgs):\n${respAll.message}")
println()
// ---- UI Extraction (custom schema) ----
runUiExtraction(providerName, providerConfig)
// --- Prompt Override ---
val prompt_req = createCompletionRequest(
providerName,
providerConfig,
modelConfig,
systemPreamble = null,
systemPromptOverride = "You are a bot named Tile Creator. Your task is to create a tile based on the user's input.",
messages=listOf(
Message(
role = Role.USER,
created = now,
content = listOf(
MessageContent.Text(
TextContent("What's your name?")
)
)
)
),
extensions=emptyList()
)
val prompt_resp = completion(prompt_req)
println("\nPrompt Override Response:\n${prompt_resp.message}")
}
/* ---------- demos ---------- */
suspend fun runCalculatorDemo(
modelConfig: ModelConfig,
providerName: String,
providerConfig: String
) {
val now = System.currentTimeMillis() / 1000
val msgs = listOf(
// same conversation you already had
Message(Role.USER, now, listOf(MessageContent.Text(TextContent("What is 7 x 6?")))),
Message(Role.ASSISTANT, now + 2, listOf(MessageContent.ToolReq(
ToolRequest(
id = "calc1",
toolCall = """
{
"status": "success",
"value": {
"name": "calculator_extension__toolname",
"arguments": { "operation": "doesnotexist", "numbers": [7,6] },
"needsApproval": false
}
}
""".trimIndent()
)))),
Message(Role.USER, now + 3, listOf(MessageContent.ToolResp(
ToolResponse(
id = "calc1",
toolResult = """
{
"status": "error",
"error": "Invalid value for operation: 'doesnotexist'. Valid values are: ['add','subtract','multiply','divide']"
}
""".trimIndent()
)))),
Message(Role.ASSISTANT, now + 4, listOf(MessageContent.ToolReq(
ToolRequest(
id = "calc1",
toolCall = """
{
"status": "success",
"value": {
"name": "calculator_extension__toolname",
"arguments": { "operation": "multiply", "numbers": [7,6] },
"needsApproval": false
}
}
""".trimIndent()
)))),
Message(Role.USER, now + 5, listOf(MessageContent.ToolResp(
ToolResponse(
id = "calc1",
toolResult = """
{
"status": "success",
"value": [ { "type": "text", "text": "42" } ]
}
""".trimIndent()
))))
)
/* one-shot prompt with error */
val reqErr = createCompletionRequest(
providerName, providerConfig, modelConfig,
"You are a helpful assistant.",
messages = listOf(msgs.first()),
extensions = listOf(calculatorExtension())
)
println("\n[${modelConfig.modelName}] Calculator (single-msg) → ${completion(reqErr).message}")
/* full conversation */
val reqAll = createCompletionRequest(
providerName, providerConfig, modelConfig,
"You are a helpful assistant.",
messages = msgs,
extensions = listOf(calculatorExtension())
)
println("[${modelConfig.modelName}] Calculator (full chat) → ${completion(reqAll).message}")
}
suspend fun runImageExample(
modelConfig: ModelConfig,
providerName: String,
providerConfig: String
) {
val imagePath = "../../crates/goose/examples/test_assets/test_image.png"
val base64Image = Base64.getEncoder().encodeToString(File(imagePath).readBytes())
val now = System.currentTimeMillis() / 1000
val msgs = listOf(
Message(Role.USER, now, listOf(
MessageContent.Text(TextContent("What is in this image?")),
MessageContent.Image(ImageContent(base64Image, "image/png"))
)),
)
val req = createCompletionRequest(
providerName, providerConfig, modelConfig,
"You are a helpful assistant. Please describe any text you see in the image.",
messages = msgs,
extensions = emptyList()
)
println("\n[${modelConfig.modelName}] Image example → ${completion(req).message}")
}
suspend fun runPromptOverride(
modelConfig: ModelConfig,
providerName: String,
providerConfig: String
) {
val now = System.currentTimeMillis() / 1000
val req = createCompletionRequest(
providerName, providerConfig, modelConfig,
systemPreamble = null,
systemPromptOverride = "You are a bot named Tile Creator. Your task is to create a tile based on the user's input.",
messages = listOf(
Message(Role.USER, now, listOf(MessageContent.Text(TextContent("What's your name?"))))
),
extensions = emptyList()
)
println("\n[${modelConfig.modelName}] Prompt override → ${completion(req).message}")
}
suspend fun runUiExtraction(providerName: String, providerConfig: String) {
val systemPrompt = "You are a UI generator AI. Convert the user input into a JSON-driven UI."
val messages = listOf(
Message(
role = Role.USER,
created = System.currentTimeMillis() / 1000,
content = listOf(
MessageContent.Text(
TextContent("Make a User Profile Form")
)
)
)
)
val schema = """{
"type": "object",
"properties": {
"type": {
"type": "string",
"enum": ["div","button","header","section","field","form"]
},
"label": { "type": "string" },
"children": {
"type": "array",
"items": { "${'$'}ref": "#" }
},
"attributes": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": { "type": "string" },
"value": { "type": "string" }
},
"required": ["name","value"],
"additionalProperties": false
}
}
},
"required": ["type","label","children","attributes"],
"additionalProperties": false
}""".trimIndent();
val schema = /* same JSON schema as before */ """
{
"type":"object",
"properties":{
"type":{"type":"string","enum":["div","button","header","section","field","form"]},
"label":{"type":"string"},
"children":{"type":"array","items":{"${'$'}ref":"#"}},
"attributes":{"type":"array","items":{"type":"object","properties":{"name":{"type":"string"},"value":{"type":"string"}},"required":["name","value"],"additionalProperties":false}}
},
"required":["type","label","children","attributes"],
"additionalProperties":false
}
""".trimIndent()
try {
val response = generateStructuredOutputs(
providerName = providerName,
providerConfig = providerConfig,
systemPrompt = systemPrompt,
messages = messages,
schema = schema
)
println("\nUI Extraction Output:\n${response}")
} catch (e: ProviderException) {
println("\nUI Extraction failed:\n${e.message}")
}
val messages = listOf(
Message(Role.USER, System.currentTimeMillis()/1000,
listOf(MessageContent.Text(TextContent("Make a User Profile Form"))))
)
val res = generateStructuredOutputs(
providerName, providerConfig,
systemPrompt = "You are a UI generator AI. Convert the user input into a JSON-driven UI.",
messages = messages,
schema = schema
)
println("\n[UI-Extraction] → $res")
}
/* ---------- entry-point ---------- */
fun main() = runBlocking {
/* --- provider setup --- */
val providerName = "databricks"
val host = System.getenv("DATABRICKS_HOST") ?: error("DATABRICKS_HOST not set")
val token = System.getenv("DATABRICKS_TOKEN") ?: error("DATABRICKS_TOKEN not set")
val providerConfig = buildProviderConfig(host, token)
println("Provider: $providerName")
println("Config : $providerConfig\n")
/* --- run demos for each model --- */
// NOTE: `claude-3-5-haiku` does NOT support images
val modelNames = listOf("kgoose-gpt-4o", "goose-claude-4-sonnet")
for (name in modelNames) {
val modelConfig = ModelConfig(name, 100000u, 0.1f, 200)
println("\n===== Running demos for model: $name =====")
runCalculatorDemo(modelConfig, providerName, providerConfig)
runImageExample(modelConfig, providerName, providerConfig)
runPromptOverride(modelConfig, providerName, providerConfig)
println("===== End demos for $name =====\n")
}
/* UI extraction is model-agnostic, so run it once */
runUiExtraction(providerName, providerConfig)
}
+16 -4
View File
@@ -833,6 +833,7 @@ internal interface UniffiLib : Library {
`systemPromptOverride`: RustBuffer.ByValue,
`messages`: RustBuffer.ByValue,
`extensions`: RustBuffer.ByValue,
`requestId`: RustBuffer.ByValue,
uniffi_out_err: UniffiRustCallStatus,
): RustBuffer.ByValue
@@ -848,6 +849,7 @@ internal interface UniffiLib : Library {
`providerName`: RustBuffer.ByValue,
`providerConfig`: RustBuffer.ByValue,
`messages`: RustBuffer.ByValue,
`requestId`: RustBuffer.ByValue,
): Long
fun uniffi_goose_llm_fn_func_generate_structured_outputs(
@@ -856,12 +858,14 @@ internal interface UniffiLib : Library {
`systemPrompt`: RustBuffer.ByValue,
`messages`: RustBuffer.ByValue,
`schema`: RustBuffer.ByValue,
`requestId`: RustBuffer.ByValue,
): Long
fun uniffi_goose_llm_fn_func_generate_tooltip(
`providerName`: RustBuffer.ByValue,
`providerConfig`: RustBuffer.ByValue,
`messages`: RustBuffer.ByValue,
`requestId`: RustBuffer.ByValue,
): Long
fun uniffi_goose_llm_fn_func_print_messages(
@@ -1101,19 +1105,19 @@ private fun uniffiCheckApiChecksums(lib: IntegrityCheckingUniffiLib) {
if (lib.uniffi_goose_llm_checksum_func_completion() != 47457.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_create_completion_request() != 50798.toShort()) {
if (lib.uniffi_goose_llm_checksum_func_create_completion_request() != 15391.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_create_tool_config() != 49910.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_generate_session_name() != 64087.toShort()) {
if (lib.uniffi_goose_llm_checksum_func_generate_session_name() != 34350.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_generate_structured_outputs() != 43426.toShort()) {
if (lib.uniffi_goose_llm_checksum_func_generate_structured_outputs() != 4576.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_generate_tooltip() != 41121.toShort()) {
if (lib.uniffi_goose_llm_checksum_func_generate_tooltip() != 36439.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_goose_llm_checksum_func_print_messages() != 30278.toShort()) {
@@ -2960,6 +2964,7 @@ fun `createCompletionRequest`(
`systemPromptOverride`: kotlin.String? = null,
`messages`: List<Message>,
`extensions`: List<ExtensionConfig>,
`requestId`: kotlin.String? = null,
): CompletionRequest =
FfiConverterTypeCompletionRequest.lift(
uniffiRustCall { _status ->
@@ -2971,6 +2976,7 @@ fun `createCompletionRequest`(
FfiConverterOptionalString.lower(`systemPromptOverride`),
FfiConverterSequenceTypeMessage.lower(`messages`),
FfiConverterSequenceTypeExtensionConfig.lower(`extensions`),
FfiConverterOptionalString.lower(`requestId`),
_status,
)
},
@@ -3003,12 +3009,14 @@ suspend fun `generateSessionName`(
`providerName`: kotlin.String,
`providerConfig`: Value,
`messages`: List<Message>,
`requestId`: kotlin.String? = null,
): kotlin.String =
uniffiRustCallAsync(
UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_session_name(
FfiConverterString.lower(`providerName`),
FfiConverterTypeValue.lower(`providerConfig`),
FfiConverterSequenceTypeMessage.lower(`messages`),
FfiConverterOptionalString.lower(`requestId`),
),
{ future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },
{ future, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_complete_rust_buffer(future, continuation) },
@@ -3031,6 +3039,7 @@ suspend fun `generateStructuredOutputs`(
`systemPrompt`: kotlin.String,
`messages`: List<Message>,
`schema`: Value,
`requestId`: kotlin.String? = null,
): ProviderExtractResponse =
uniffiRustCallAsync(
UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_structured_outputs(
@@ -3039,6 +3048,7 @@ suspend fun `generateStructuredOutputs`(
FfiConverterString.lower(`systemPrompt`),
FfiConverterSequenceTypeMessage.lower(`messages`),
FfiConverterTypeValue.lower(`schema`),
FfiConverterOptionalString.lower(`requestId`),
),
{ future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },
{ future, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_complete_rust_buffer(future, continuation) },
@@ -3059,12 +3069,14 @@ suspend fun `generateTooltip`(
`providerName`: kotlin.String,
`providerConfig`: Value,
`messages`: List<Message>,
`requestId`: kotlin.String? = null,
): kotlin.String =
uniffiRustCallAsync(
UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_tooltip(
FfiConverterString.lower(`providerName`),
FfiConverterTypeValue.lower(`providerConfig`),
FfiConverterSequenceTypeMessage.lower(`messages`),
FfiConverterOptionalString.lower(`requestId`),
),
{ future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },
{ future, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_complete_rust_buffer(future, continuation) },