chore: remove ffi libraries and related code (#3699)

This commit is contained in:
Salman Mohammed
2025-07-29 09:21:08 -04:00
committed by GitHub
parent af2ca33bc1
commit 4ac5fe4f8d
52 changed files with 15 additions and 14121 deletions
Generated
+15 -413
View File
@@ -430,48 +430,6 @@ dependencies = [
"regex-syntax 0.8.5",
]
[[package]]
name = "askama"
version = "0.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d4744ed2eef2645831b441d8f5459689ade2ab27c854488fbab1fbe94fce1a7"
dependencies = [
"askama_derive",
"itoa",
"percent-encoding",
"serde",
"serde_json",
]
[[package]]
name = "askama_derive"
version = "0.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d661e0f57be36a5c14c48f78d09011e67e0cb618f269cca9f2fd8d15b68c46ac"
dependencies = [
"askama_parser",
"basic-toml",
"memchr",
"proc-macro2",
"quote",
"rustc-hash 2.1.1",
"serde",
"serde_derive",
"syn 2.0.99",
]
[[package]]
name = "askama_parser"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cf315ce6524c857bb129ff794935cf6d42c82a6cff60526fe2a63593de4d0d4f"
dependencies = [
"memchr",
"serde",
"serde_derive",
"winnow",
]
[[package]]
name = "assert-json-diff"
version = "2.0.2"
@@ -482,19 +440,6 @@ dependencies = [
"serde_json",
]
[[package]]
name = "async-compat"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7bab94bde396a3f7b4962e396fdad640e241ed797d4d8d77fc8c237d14c58fc0"
dependencies = [
"futures-core",
"futures-io",
"once_cell",
"pin-project-lite",
"tokio",
]
[[package]]
name = "async-compression"
version = "0.4.20"
@@ -614,17 +559,6 @@ version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
[[package]]
name = "atty"
version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8"
dependencies = [
"hermit-abi 0.1.19",
"libc",
"winapi",
]
[[package]]
name = "autocfg"
version = "1.4.0"
@@ -1191,15 +1125,6 @@ dependencies = [
"vsimd",
]
[[package]]
name = "basic-toml"
version = "0.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba62675e8242a4c4e806d12f11d136e626e6c8361d6b829310732241652a178a"
dependencies = [
"serde",
]
[[package]]
name = "bat"
version = "0.24.0"
@@ -1210,7 +1135,7 @@ dependencies = [
"bincode",
"bugreport",
"bytesize",
"clap 4.5.31",
"clap",
"clircle",
"console",
"content_inspector",
@@ -1516,20 +1441,6 @@ dependencies = [
"serde_json",
]
[[package]]
name = "cargo_metadata"
version = "0.19.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dd5eb614ed4c27c5d706420e4320fbe3216ab31fa1c33cd8246ac36dae4479ba"
dependencies = [
"camino",
"cargo-platform",
"semver",
"serde",
"serde_json",
"thiserror 2.0.12",
]
[[package]]
name = "cast"
version = "0.3.0"
@@ -1545,25 +1456,6 @@ dependencies = [
"cipher",
]
[[package]]
name = "cbindgen"
version = "0.24.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b922faaf31122819ec80c4047cc684c6979a087366c069611e33649bf98e18d"
dependencies = [
"clap 3.2.25",
"heck 0.4.1",
"indexmap 1.9.3",
"log",
"proc-macro2",
"quote",
"serde",
"serde_json",
"syn 1.0.109",
"tempfile",
"toml 0.5.11",
]
[[package]]
name = "cc"
version = "1.2.16"
@@ -1714,21 +1606,6 @@ dependencies = [
"libloading",
]
[[package]]
name = "clap"
version = "3.2.25"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ea181bf566f71cb9a5d17a59e1871af638180a18fb0035c92ae62b705207123"
dependencies = [
"atty",
"bitflags 1.3.2",
"clap_lex 0.2.4",
"indexmap 1.9.3",
"strsim 0.10.0",
"termcolor",
"textwrap",
]
[[package]]
name = "clap"
version = "4.5.31"
@@ -1747,8 +1624,8 @@ checksum = "5589e0cba072e0f3d23791efac0fd8627b49c829c196a492e88168e6a669d863"
dependencies = [
"anstream",
"anstyle",
"clap_lex 0.7.4",
"strsim 0.11.1",
"clap_lex",
"strsim",
"terminal_size",
]
@@ -1764,15 +1641,6 @@ dependencies = [
"syn 2.0.99",
]
[[package]]
name = "clap_lex"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2850f2f5a82cbf437dd5af4d49848fbdfc27c157c3d010345776f952765261c5"
dependencies = [
"os_str_bytes",
]
[[package]]
name = "clap_lex"
version = "0.7.4"
@@ -1788,7 +1656,7 @@ dependencies = [
"console",
"indicatif",
"once_cell",
"strsim 0.11.1",
"strsim",
"textwrap",
"zeroize",
]
@@ -1879,7 +1747,7 @@ dependencies = [
"rust-ini",
"serde",
"serde_json",
"toml 0.8.20",
"toml",
"yaml-rust2",
]
@@ -2046,7 +1914,7 @@ dependencies = [
"anes",
"cast",
"ciborium",
"clap 4.5.31",
"clap",
"criterion-plot",
"is-terminal",
"itertools 0.10.5",
@@ -2191,7 +2059,7 @@ dependencies = [
"ident_case",
"proc-macro2",
"quote",
"strsim 0.11.1",
"strsim",
"syn 2.0.99",
]
@@ -3090,15 +2958,6 @@ version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c2141d6d6c8512188a7891b4b01590a45f6dac67afb4f255c4124dbb86d4eaa"
[[package]]
name = "fs-err"
version = "2.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88a41f105fe1d5b6b34b2055e3dc59bb79b46b48b2040b9e6c7b4b5de097aa41"
dependencies = [
"autocfg",
]
[[package]]
name = "fs2"
version = "0.4.3"
@@ -3341,17 +3200,6 @@ dependencies = [
"regex-syntax 0.8.5",
]
[[package]]
name = "goblin"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b363a30c165f666402fe6a3024d3bec7ebc898f96a4a23bd1c99f8dbf3f4f47"
dependencies = [
"log",
"plain",
"scroll",
]
[[package]]
name = "google-apis-common"
version = "7.0.0"
@@ -3461,7 +3309,7 @@ dependencies = [
"fs2",
"futures",
"include_dir",
"indoc 2.0.6",
"indoc",
"jsonschema",
"jsonwebtoken",
"keyring",
@@ -3521,7 +3369,7 @@ dependencies = [
"serde",
"serde_json",
"tokio",
"toml 0.8.20",
"toml",
"tracing",
"tracing-subscriber",
"winapi",
@@ -3538,7 +3386,7 @@ dependencies = [
"bat",
"bytes",
"chrono",
"clap 4.5.31",
"clap",
"cliclack",
"console",
"dirs 5.0.1",
@@ -3581,50 +3429,6 @@ dependencies = [
"winapi",
]
[[package]]
name = "goose-ffi"
version = "1.1.0"
dependencies = [
"cbindgen",
"futures",
"goose",
"libc",
"once_cell",
"serde",
"serde_json",
"tokio",
]
[[package]]
name = "goose-llm"
version = "1.1.0"
dependencies = [
"anyhow",
"async-trait",
"base64 0.21.7",
"chrono",
"criterion",
"ctor",
"dotenvy",
"goose",
"include_dir",
"indoc 1.0.9",
"lazy_static",
"minijinja",
"once_cell",
"regex",
"reqwest 0.12.12",
"serde",
"serde_json",
"smallvec",
"tempfile",
"thiserror 1.0.69",
"tokio",
"tracing",
"uniffi",
"url",
]
[[package]]
name = "goose-mcp"
version = "1.1.0"
@@ -3645,7 +3449,7 @@ dependencies = [
"ignore",
"image 0.24.9",
"include_dir",
"indoc 2.0.6",
"indoc",
"keyring",
"kill_tree",
"lazy_static",
@@ -3688,7 +3492,7 @@ dependencies = [
"base64 0.21.7",
"bytes",
"chrono",
"clap 4.5.31",
"clap",
"config",
"dirs 6.0.0",
"etcetera",
@@ -3828,15 +3632,6 @@ version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "hermit-abi"
version = "0.1.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33"
dependencies = [
"libc",
]
[[package]]
name = "hermit-abi"
version = "0.3.9"
@@ -4389,12 +4184,6 @@ dependencies = [
"web-time",
]
[[package]]
name = "indoc"
version = "1.0.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa799dd5ed20a7e349f3b4639aa80d74549c81716d9ec4f994c9b5815598306"
[[package]]
name = "indoc"
version = "2.0.6"
@@ -6092,12 +5881,6 @@ dependencies = [
"hashbrown 0.14.5",
]
[[package]]
name = "os_str_bytes"
version = "6.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e2355d85b9a3786f481747ced0e0ff2ba35213a1f9bd406ed906554d7af805a1"
[[package]]
name = "outref"
version = "0.5.2"
@@ -6293,7 +6076,7 @@ version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5"
dependencies = [
"siphasher 1.0.1",
"siphasher",
]
[[package]]
@@ -6334,12 +6117,6 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
[[package]]
name = "plain"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
[[package]]
name = "plist"
version = "1.7.0"
@@ -7481,26 +7258,6 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "scroll"
version = "0.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ab8598aa408498679922eff7fa985c25d58a90771bd6be794434c5277eab1a6"
dependencies = [
"scroll_derive",
]
[[package]]
name = "scroll_derive"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1783eabc414609e28a5ba76aee5ddd52199f7107a0b24c2e9746a1ecc34a683d"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.99",
]
[[package]]
name = "sct"
version = "0.7.1"
@@ -7822,12 +7579,6 @@ dependencies = [
"time",
]
[[package]]
name = "siphasher"
version = "0.3.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d"
[[package]]
name = "siphasher"
version = "1.0.1"
@@ -7841,7 +7592,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "16d23b015676c90a0f01c197bfdc786c20342c73a0afdda9025adb0bc42940a8"
dependencies = [
"bytecount",
"cargo_metadata 0.14.2",
"cargo_metadata",
"error-chain",
"glob",
"pulldown-cmark",
@@ -7872,9 +7623,6 @@ name = "smallvec"
version = "1.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd"
dependencies = [
"serde",
]
[[package]]
name = "smawk"
@@ -7969,12 +7717,6 @@ version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e51f1e89f093f99e7432c491c382b88a6860a5adbe6bf02574bf0a08efff1978"
[[package]]
name = "strsim"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623"
[[package]]
name = "strsim"
version = "0.11.1"
@@ -8137,7 +7879,7 @@ dependencies = [
"cfg-expr",
"heck 0.5.0",
"pkg-config",
"toml 0.8.20",
"toml",
"version-compare",
]
@@ -8680,15 +8422,6 @@ dependencies = [
"tokio",
]
[[package]]
name = "toml"
version = "0.5.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234"
dependencies = [
"serde",
]
[[package]]
name = "toml"
version = "0.8.20"
@@ -8996,128 +8729,6 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"
[[package]]
name = "uniffi"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4dcd1d240101ba3b9d7532ae86d9cb64d9a7ff63e13a2b7b9e94a32a601d8233"
dependencies = [
"anyhow",
"camino",
"cargo_metadata 0.19.2",
"clap 4.5.31",
"uniffi_bindgen",
"uniffi_core",
"uniffi_macros",
"uniffi_pipeline",
]
[[package]]
name = "uniffi_bindgen"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d0525f06d749ea80d8049dc0bb038bb87941e3d909eefa76b6f0a5589b59ac5"
dependencies = [
"anyhow",
"askama",
"camino",
"cargo_metadata 0.19.2",
"fs-err",
"glob",
"goblin",
"heck 0.5.0",
"indexmap 2.7.1",
"once_cell",
"serde",
"tempfile",
"textwrap",
"toml 0.5.11",
"uniffi_internal_macros",
"uniffi_meta",
"uniffi_pipeline",
"uniffi_udl",
]
[[package]]
name = "uniffi_core"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3fa8eb4d825b4ed095cb13483cba6927c3002b9eb603cef9b7688758cc3772e"
dependencies = [
"anyhow",
"async-compat",
"bytes",
"once_cell",
"static_assertions",
]
[[package]]
name = "uniffi_internal_macros"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "83b547d69d699e52f2129fde4b57ae0d00b5216e59ed5b56097c95c86ba06095"
dependencies = [
"anyhow",
"indexmap 2.7.1",
"proc-macro2",
"quote",
"syn 2.0.99",
]
[[package]]
name = "uniffi_macros"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "00f1de72edc8cb9201c7d650e3678840d143e4499004571aac49e6cb1b17da43"
dependencies = [
"camino",
"fs-err",
"once_cell",
"proc-macro2",
"quote",
"serde",
"syn 2.0.99",
"toml 0.5.11",
"uniffi_meta",
]
[[package]]
name = "uniffi_meta"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3acc9204632f6a555b2cba7c8852c5523bc1aa5f3eff605c64af5054ea28b72e"
dependencies = [
"anyhow",
"siphasher 0.3.11",
"uniffi_internal_macros",
"uniffi_pipeline",
]
[[package]]
name = "uniffi_pipeline"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54b5336a9a925b358183837d31541d12590b7fcec373256d3770de02dff24c69"
dependencies = [
"anyhow",
"heck 0.5.0",
"indexmap 2.7.1",
"tempfile",
"uniffi_internal_macros",
]
[[package]]
name = "uniffi_udl"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f95e73373d85f04736bc51997d3e6855721144ec4384cae9ca8513c80615e129"
dependencies = [
"anyhow",
"textwrap",
"uniffi_meta",
"weedle2",
]
[[package]]
name = "unsafe-libyaml"
version = "0.2.11"
@@ -9452,15 +9063,6 @@ dependencies = [
"rustls-pki-types",
]
[[package]]
name = "weedle2"
version = "5.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "998d2c24ec099a87daf9467808859f9d82b61f1d9c9701251aea037f514eae0e"
dependencies = [
"nom",
]
[[package]]
name = "weezl"
version = "0.1.8"
-20
View File
@@ -444,23 +444,3 @@ win-total-rls *allparam:
just win-bld-rls{{allparam}}
just win-run-rls
### Build and run the Kotlin example with
### auto-generated bindings for goose-llm
kotlin-example:
# Build Rust dylib and generate Kotlin bindings
cargo build -p goose-llm
cargo run --features=uniffi/cli --bin uniffi-bindgen generate \
--library ./target/debug/libgoose_llm.dylib --language kotlin --out-dir bindings/kotlin
# Compile and run the Kotlin example
cd bindings/kotlin/ && kotlinc \
example/Usage.kt \
uniffi/goose_llm/goose_llm.kt \
-classpath "libs/kotlin-stdlib-1.9.0.jar:libs/kotlinx-coroutines-core-jvm-1.7.3.jar:libs/jna-5.13.0.jar" \
-include-runtime \
-d example.jar
cd bindings/kotlin/ && java \
-Djna.library.path=$HOME/Development/goose/target/debug \
-classpath "example.jar:libs/kotlin-stdlib-1.9.0.jar:libs/kotlinx-coroutines-core-jvm-1.7.3.jar:libs/jna-5.13.0.jar" \
UsageKt
-115
View File
@@ -1,115 +0,0 @@
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))
}
-228
View File
@@ -1,228 +0,0 @@
import java.io.File
import java.util.Base64
import kotlinx.coroutines.runBlocking
import uniffi.goose_llm.*
/* ---------- shared helpers ---------- */
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",
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"
}
}
}
""".trimIndent(),
approvalMode = ToolApprovalMode.AUTO
)
return ExtensionConfig(
name = "calculator_extension",
instructions = "This extension provides a calculator tool.",
tools = listOf(calculatorTool)
)
}
/* ---------- 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 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()
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)
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
-133
View File
@@ -1,133 +0,0 @@
import asyncio
import os
import time
from goose_llm import (
Message, MessageContent, TextContent, ToolRequest, ToolResponse,
Role, ModelConfig, ToolApprovalMode,
create_tool_config, ExtensionConfig,
generate_session_name, generate_tooltip,
create_completion_request, completion
)
async def main():
now = int(time.time())
# 1) User sends a plain-text prompt
messages = [
Message(
role=Role.USER,
created=now,
content=[MessageContent.TEXT(TextContent(text="What is 7 x 6?"))]
),
# 2) Assistant makes a tool request
Message(
role=Role.ASSISTANT,
created=now + 2,
content=[MessageContent.TOOL_REQ(ToolRequest(
id="calc1",
tool_call="""
{
"status": "success",
"value": {
"name": "calculator_extension__toolname",
"arguments": {
"operation": "multiply",
"numbers": [7, 6]
},
"needsApproval": false
}
}
"""
))]
),
# 3) User sends tool result
Message(
role=Role.USER,
created=now + 3,
content=[MessageContent.TOOL_RESP(ToolResponse(
id="calc1",
tool_result="""
{
"status": "success",
"value": [
{"type": "text", "text": "42"}
]
}
"""
))]
)
]
provider_name = "databricks"
provider_config = f'''{{
"host": "{os.environ.get("DATABRICKS_HOST")}",
"token": "{os.environ.get("DATABRICKS_TOKEN")}"
}}'''
print(f"Provider Name: {provider_name}")
print(f"Provider Config: {provider_config}")
session_name = await generate_session_name(provider_name, provider_config, messages)
print(f"\nSession Name: {session_name}")
tooltip = await generate_tooltip(provider_name, provider_config, messages)
print(f"\nTooltip: {tooltip}")
model_config = ModelConfig(
model_name="goose-gpt-4-1",
max_tokens=500,
temperature=0.1,
context_limit=4096,
)
calculator_tool = create_tool_config(
name="calculator",
description="Perform basic arithmetic operations",
input_schema="""
{
"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"
}
}
}
""",
approval_mode=ToolApprovalMode.AUTO
)
calculator_extension = ExtensionConfig(
name="calculator_extension",
instructions="This extension provides a calculator tool.",
tools=[calculator_tool]
)
system_preamble = "You are a helpful assistant."
extensions = [calculator_extension]
req = create_completion_request(
provider_name,
provider_config,
model_config,
system_preamble,
messages,
extensions
)
resp = await completion(req)
print(f"\nCompletion Response:\n{resp.message}")
print(f"Msg content: {resp.message.content[0][0]}")
if __name__ == "__main__":
asyncio.run(main())
-28
View File
@@ -1,28 +0,0 @@
[package]
name = "goose-ffi"
build = "build.rs"
version.workspace = true
edition.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[lints]
workspace = true
[lib]
name = "goose_ffi"
crate-type = ["cdylib"]
[dependencies]
goose = { path = "../goose" }
futures = "0.3"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
tokio = { version = "1", features = ["full"] }
libc = "0.2"
once_cell = "1.18"
[build-dependencies]
cbindgen = "0.24.0"
-127
View File
@@ -1,127 +0,0 @@
# Goose FFI
Foreign Function Interface (FFI) for the Goose AI agent framework, allowing integration with other programming languages.
## Overview
The Goose FFI library provides C-compatible bindings for the Goose AI agent framework, enabling you to:
- Create and manage Goose agents from any language with C FFI support
- Configure and use the Databricks AI provider for now but is extensible to other providers as needed
- Send messages to agents and receive responses
## Building
To build the FFI library, you'll need Rust and Cargo installed. Then run:
```bash
# Build the library in debug mode
cargo build --package goose-ffi
# Build the library in release mode (recommended for production)
cargo build --release --package goose-ffi
```
This will generate a dynamic library (.so, .dll, or .dylib depending on your platform) in the `target` directory, and automatically generate the C header file in the `include` directory.
You can also build cross-platform binaries using cross command. For example to build for linux x86_64 architecture from Mac would require running
```bash
CROSS_BUILD_OPTS="--platform linux/amd64 --no-cache" CROSS_CONTAINER_OPTS="--platform linux/amd64" cross build -p goose-ffi --release --target x86_64-unknown-linux-gnu --no-default-features
```
Note that this works only for gnu linux as it requires glibc.
## Generated C Header
The library uses cbindgen to automatically generate a C header file (`goose_ffi.h`) during the build process. This header contains all the necessary types and function declarations to use the library from C or any language with C FFI support.
## Examples
The FFI library includes examples in multiple languages to demonstrate how to use it.
### Python Example
The `examples/goose_agent.py` demonstrates using the FFI library from Python with ctypes. It shows:
1. How to create a proper Python wrapper around the Goose FFI interface
2. Loading the shared library dynamically based on platform
3. Setting up C-compatible structures
4. Creating an object-oriented API for easier use
Note: Tool callback functionality shown in earlier versions is not currently available and will be implemented in a future release.
To run the Python example:
```bash
# First, build the FFI library
cargo build --release --package goose-ffi
# Then set the environment variables & run the example
DATABRICKS_HOST=... DATABRICKS_API_KEY=... python crates/goose-ffi/examples/goose_agent.py
```
You need to have Python 3.6+ installed with the `ctypes` module (included in standard library).
```
> Tell me about the Eiffel Tower
```
The agent will respond with information about the Eiffel Tower.
## Using from Other Languages
The Goose FFI library can be used from many programming languages with C FFI support, including:
- Python (via ctypes or cffi)
- JavaScript/Node.js (via node-ffi)
- Ruby (via fiddle)
- C#/.NET (via P/Invoke)
- Go (via cgo)
- Java / Kotlin (via JNA or JNI)
Check the documentation for FFI support in your language of choice for details on how to load and use a C shared library.
## Provider Configuration
The FFI interface uses a provider type enumeration to specify which AI provider to use:
```c
// C enum (defined in examples/simple_agent.c)
typedef enum {
PROVIDER_DATABRICKS = 0, // Databricks AI provider
} ProviderType;
```
```python
# Python enum (defined in examples/goose_agent.py)
class ProviderType(IntEnum):
DATABRICKS = 0 # Databricks AI provider
```
Currently, only the Databricks provider (provider_type = 0) is supported. If you attempt to use any other provider type, an error will be returned.
### Environment-based Configuration
The library supports configuration via environment variables, which makes it easier to use in containerized or CI/CD environments without hardcoding credentials:
#### Databricks Provider (type = 0)
```
DATABRICKS_API_KEY=dapi... # Databricks API key
DATABRICKS_HOST=... # Databricks host URL (e.g., "https://your-workspace.cloud.databricks.com")
```
These environment variables will be used automatically if you don't provide the corresponding parameters when creating an agent.
## Thread Safety
The FFI library is designed to be thread-safe. Each agent instance is independent, and tools callbacks are handled in a thread-safe manner. However, the same agent instance should not be used from multiple threads simultaneously without external synchronization.
## Error Handling
Functions that can fail return either null pointers or special result structures that indicate success or failure. Always check return values and clean up resources using the appropriate free functions.
## Memory Management
The FFI interface handles memory allocation and deallocation. Use the provided free functions (like `goose_free_string` and `goose_free_async_result`) to release resources when you're done with them.
-48
View File
@@ -1,48 +0,0 @@
use std::env;
use std::path::PathBuf;
fn main() {
let crate_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
let config = cbindgen::Config {
language: cbindgen::Language::C,
documentation: true,
header: Some(
r#"
#ifndef GOOSE_FFI_H
#define GOOSE_FFI_H
/* Goose FFI - C interface for the Goose AI agent framework */
"#
.trim_start()
.to_string(),
),
trailer: Some("#endif // GOOSE_FFI_H".to_string()),
includes: vec![],
sys_includes: vec!["stdint.h".to_string(), "stdbool.h".to_string()],
export: cbindgen::ExportConfig {
prefix: Some("goose_".to_string()),
..Default::default()
},
documentation_style: cbindgen::DocumentationStyle::C,
enumeration: cbindgen::EnumConfig {
prefix_with_name: true,
derive_helper_methods: true,
..Default::default()
},
..Default::default()
};
let bindings = cbindgen::Builder::new()
.with_crate(&crate_dir)
.with_config(config)
.generate()
.expect("Unable to generate bindings");
let out_path = PathBuf::from(&crate_dir).join("include");
std::fs::create_dir_all(&out_path).expect("Failed to create include directory");
bindings.write_to_file(out_path.join("goose_ffi.h"));
println!("cargo:rerun-if-changed=src/lib.rs");
println!("cargo:rerun-if-changed=build.rs");
}
-116
View File
@@ -1,116 +0,0 @@
#!/usr/bin/env python3
"""
Python example for using the Goose FFI interface.
This example demonstrates how to:
1. Load the Goose FFI library
2. Create an agent with a provider
3. Add a tool extension
4. Send messages to the agent
5. Handle tool calls and responses
"""
import ctypes
import os
import platform
from ctypes import c_char_p, c_bool, c_uint32, c_void_p, Structure, POINTER
class ProviderType:
DATABRICKS = 0
# Platform-specific dynamic lib name
if platform.system() == "Darwin":
LIB_NAME = "libgoose_ffi.dylib"
elif platform.system() == "Linux":
LIB_NAME = "libgoose_ffi.so"
elif platform.system() == "Windows":
LIB_NAME = "goose_ffi.dll"
else:
raise RuntimeError("Unsupported platform")
# Adjust to your actual build output directory
LIB_PATH = os.path.join(os.path.dirname(__file__), "../../..", "target", "debug", LIB_NAME)
# Load library
goose = ctypes.CDLL(LIB_PATH)
# Forward declaration for goose_Agent
class goose_Agent(Structure):
pass
# Agent pointer type
goose_AgentPtr = POINTER(goose_Agent)
# C struct mappings
class ProviderConfig(Structure):
_fields_ = [
("provider_type", c_uint32),
("api_key", c_char_p),
("model_name", c_char_p),
("host", c_char_p),
]
class AsyncResult(Structure):
_fields_ = [
("succeeded", c_bool),
("error_message", c_char_p),
]
# Function signatures
goose.goose_agent_new.argtypes = [POINTER(ProviderConfig)]
goose.goose_agent_new.restype = goose_AgentPtr
goose.goose_agent_free.argtypes = [goose_AgentPtr]
goose.goose_agent_free.restype = None
goose.goose_agent_send_message.argtypes = [goose_AgentPtr, c_char_p]
goose.goose_agent_send_message.restype = c_void_p
goose.goose_free_string.argtypes = [c_void_p]
goose.goose_free_string.restype = None
goose.goose_free_async_result.argtypes = [POINTER(AsyncResult)]
goose.goose_free_async_result.restype = None
class GooseAgent:
def __init__(self, provider_type=ProviderType.DATABRICKS, api_key=None, model_name=None, host=None):
self.config = ProviderConfig(
provider_type=provider_type,
api_key=api_key.encode("utf-8") if api_key else None,
model_name=model_name.encode("utf-8") if model_name else None,
host=host.encode("utf-8") if host else None,
)
self.agent = goose.goose_agent_new(ctypes.byref(self.config))
if not self.agent:
raise RuntimeError("Failed to create Goose agent")
def __del__(self):
if getattr(self, "agent", None):
goose.goose_agent_free(self.agent)
def send_message(self, message: str) -> str:
msg = message.encode("utf-8")
response_ptr = goose.goose_agent_send_message(self.agent, msg)
if not response_ptr:
return "Error or NULL response from agent"
response = ctypes.string_at(response_ptr).decode("utf-8")
# Free the string using the proper C function provided by the library
# This correctly releases memory allocated by the Rust side
goose.goose_free_string(response_ptr)
return response
def main():
api_key = os.getenv("DATABRICKS_API_KEY")
host = os.getenv("DATABRICKS_HOST")
agent = GooseAgent(api_key=api_key, model_name="claude-3-7-sonnet", host=host)
print("Type a message (or 'quit' to exit):")
while True:
user_input = input("> ")
if user_input.lower() in ("quit", "exit"):
break
reply = agent.send_message(user_input)
print(f"Agent: {reply}\n")
if __name__ == "__main__":
main()
-145
View File
@@ -1,145 +0,0 @@
#ifndef GOOSE_FFI_H
#define GOOSE_FFI_H
/* Goose FFI - C interface for the Goose AI agent framework */
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
/*
Provider Type enumeration
Currently only Databricks is supported
*/
enum goose_ProviderType {
/*
Databricks AI provider
*/
goose_ProviderType_Databricks = 0,
};
typedef uint32_t goose_ProviderType;
/*
Result type for async operations
- succeeded: true if the operation succeeded, false otherwise
- error_message: Error message if succeeded is false, NULL otherwise
*/
typedef struct goose_AsyncResult {
bool succeeded;
char *error_message;
} goose_AsyncResult;
/*
Pointer type for the agent
*/
typedef goose_Agent *goose_AgentPtr;
/*
Provider configuration used to initialize an AI provider
- provider_type: Provider type (0 = Databricks, other values will produce an error)
- api_key: Provider API key (null for default from environment variables)
- model_name: Model name to use (null for provider default)
- host: Provider host URL (null for default from environment variables)
*/
typedef struct goose_ProviderConfigFFI {
goose_ProviderType provider_type;
const char *api_key;
const char *model_name;
const char *host;
} goose_ProviderConfigFFI;
/*
Free an async result structure
This function frees the memory allocated for an AsyncResult structure,
including any error message it contains.
# Safety
The result pointer must be a valid pointer returned by a goose FFI function,
or NULL.
*/
void goose_free_async_result(struct goose_AsyncResult *result);
/*
Create a new agent with the given provider configuration
# Parameters
- config: Provider configuration
# Returns
A new agent pointer, or a null pointer if creation failed
# Safety
The config pointer must be valid or NULL. The resulting agent must be freed
with goose_agent_free when no longer needed.
*/
goose_AgentPtr goose_agent_new(const struct goose_ProviderConfigFFI *config);
/*
Free an agent
This function frees the memory allocated for an agent.
# Parameters
- agent_ptr: Agent pointer returned by goose_agent_new
# Safety
The agent_ptr must be a valid pointer returned by goose_agent_new,
or have a null internal pointer. The agent_ptr must not be used after
calling this function.
*/
void goose_agent_free(goose_AgentPtr agent_ptr);
/*
Send a message to the agent and get the response
This function sends a message to the agent and returns the response.
Tool handling is not yet supported and will be implemented in a future commit
so this may change significantly
# Parameters
- agent_ptr: Agent pointer
- message: Message to send
# Returns
A C string with the agent's response, or NULL on error.
This string must be freed with goose_free_string when no longer needed.
# Safety
The agent_ptr must be a valid pointer returned by goose_agent_new.
The message must be a valid C string.
*/
char *goose_agent_send_message(goose_AgentPtr agent_ptr, const char *message);
/*
Free a string allocated by goose FFI functions
This function frees memory allocated for strings returned by goose FFI functions.
# Parameters
- s: String to free
# Safety
The string must have been allocated by a goose FFI function, or be NULL.
The string must not be used after calling this function.
*/
void goose_free_string(char *s);
#endif // GOOSE_FFI_H
-311
View File
@@ -1,311 +0,0 @@
use std::ffi::{c_char, CStr, CString};
use std::ptr;
use std::sync::Arc;
use futures::StreamExt;
use goose::agents::{Agent, AgentEvent};
use goose::message::Message;
use goose::model::ModelConfig;
use goose::providers::databricks::DatabricksProvider;
use once_cell::sync::OnceCell;
use tokio::runtime::Runtime;
// This class is in alpha and not yet ready for production use
// and the API is not yet stable. Use at your own risk.
// Thread-safe global runtime
static RUNTIME: OnceCell<Runtime> = OnceCell::new();
// Get or initialize the global runtime
fn get_runtime() -> &'static Runtime {
RUNTIME.get_or_init(|| {
// Runtime with all features enabled
Runtime::new().expect("Failed to create Tokio runtime")
})
}
/// Pointer type for the agent
pub type AgentPtr = *mut Agent;
/// Provider Type enumeration
/// Currently only Databricks is supported
#[repr(u32)]
#[derive(Debug, Clone, Copy)]
pub enum ProviderType {
/// Databricks AI provider
Databricks = 0,
}
/// Provider configuration used to initialize an AI provider
///
/// - provider_type: Provider type (0 = Databricks, other values will produce an error)
/// - api_key: Provider API key (null for default from environment variables)
/// - model_name: Model name to use (null for provider default)
/// - host: Provider host URL (null for default from environment variables)
#[repr(C)]
pub struct ProviderConfigFFI {
pub provider_type: ProviderType,
pub api_key: *const c_char,
pub model_name: *const c_char,
pub host: *const c_char,
}
// Extension configuration will be implemented in a future commit
/// Role enum for message participants
#[repr(u32)]
#[derive(Debug, Clone, Copy)]
pub enum MessageRole {
/// User message role
User = 0,
/// Assistant message role
Assistant = 1,
/// System message role
System = 2,
}
/// Message structure for agent interactions
///
/// - role: Message role (User, Assistant, or System)
/// - content: Text content of the message
#[repr(C)]
pub struct MessageFFI {
pub role: MessageRole,
pub content: *const c_char,
}
// Tool callbacks will be implemented in a future commit
/// Result type for async operations
///
/// - succeeded: true if the operation succeeded, false otherwise
/// - error_message: Error message if succeeded is false, NULL otherwise
#[repr(C)]
pub struct AsyncResult {
pub succeeded: bool,
pub error_message: *mut c_char,
}
/// Free an async result structure
///
/// This function frees the memory allocated for an AsyncResult structure,
/// including any error message it contains.
///
/// # Safety
///
/// The result pointer must be a valid pointer returned by a goose FFI function,
/// or NULL.
#[no_mangle]
pub unsafe extern "C" fn goose_free_async_result(result: *mut AsyncResult) {
if !result.is_null() {
let result = &mut *result;
if !result.error_message.is_null() {
let _ = CString::from_raw(result.error_message);
}
let _ = Box::from_raw(result);
}
}
/// Create a new agent with the given provider configuration
///
/// # Parameters
///
/// - config: Provider configuration
///
/// # Returns
///
/// A new agent pointer, or a null pointer if creation failed
///
/// # Safety
///
/// The config pointer must be valid or NULL. The resulting agent must be freed
/// with goose_agent_free when no longer needed.
#[no_mangle]
pub unsafe extern "C" fn goose_agent_new(config: *const ProviderConfigFFI) -> AgentPtr {
// Check for null pointer
if config.is_null() {
eprintln!("Error: config pointer is null");
return ptr::null_mut();
}
let config = &*config;
// We currently only support Databricks provider
// This match ensures future compiler errors if new provider types are added without handling
match config.provider_type {
ProviderType::Databricks => (), // Databricks provider is supported
}
// Get api_key from config or environment
let api_key = if !config.api_key.is_null() {
CStr::from_ptr(config.api_key).to_string_lossy().to_string()
} else {
match std::env::var("DATABRICKS_API_KEY") {
Ok(key) => key,
Err(_) => {
eprintln!("Error: api_key not provided and DATABRICKS_API_KEY environment variable not set");
return ptr::null_mut();
}
}
};
// Check and get required model_name (no env fallback for model)
if config.model_name.is_null() {
eprintln!("Error: model_name is required but was null");
return ptr::null_mut();
}
let model_name = CStr::from_ptr(config.model_name)
.to_string_lossy()
.to_string();
// Get host from config or environment
let host = if !config.host.is_null() {
CStr::from_ptr(config.host).to_string_lossy().to_string()
} else {
match std::env::var("DATABRICKS_HOST") {
Ok(url) => url,
Err(_) => {
eprintln!(
"Error: host not provided and DATABRICKS_HOST environment variable not set"
);
return ptr::null_mut();
}
}
};
// Create model config with model name
let model_config = ModelConfig::new(model_name);
// Create Databricks provider with required parameters
match DatabricksProvider::from_params(host, api_key, model_config) {
Ok(provider) => {
let agent = Agent::new();
get_runtime().block_on(async {
let _ = agent.update_provider(Arc::new(provider)).await;
});
Box::into_raw(Box::new(agent))
}
Err(e) => {
eprintln!("Error creating Databricks provider: {:?}", e);
ptr::null_mut()
}
}
}
/// Free an agent
///
/// This function frees the memory allocated for an agent.
///
/// # Parameters
///
/// - agent_ptr: Agent pointer returned by goose_agent_new
///
/// # Safety
///
/// The agent_ptr must be a valid pointer returned by goose_agent_new,
/// or have a null internal pointer. The agent_ptr must not be used after
/// calling this function.
#[no_mangle]
pub unsafe extern "C" fn goose_agent_free(agent_ptr: AgentPtr) {
if !agent_ptr.is_null() {
let _ = Box::from_raw(agent_ptr);
}
}
/// Send a message to the agent and get the response
///
/// This function sends a message to the agent and returns the response.
/// Tool handling is not yet supported and will be implemented in a future commit
/// so this may change significantly
///
/// # Parameters
///
/// - agent_ptr: Agent pointer
/// - message: Message to send
///
/// # Returns
///
/// A C string with the agent's response, or NULL on error.
/// This string must be freed with goose_free_string when no longer needed.
///
/// # Safety
///
/// The agent_ptr must be a valid pointer returned by goose_agent_new.
/// The message must be a valid C string.
#[no_mangle]
pub unsafe extern "C" fn goose_agent_send_message(
agent_ptr: AgentPtr,
message: *const c_char,
) -> *mut c_char {
if agent_ptr.is_null() || message.is_null() {
return ptr::null_mut();
}
let agent = &mut *agent_ptr;
let message = CStr::from_ptr(message).to_string_lossy().to_string();
let messages = vec![Message::user().with_text(&message)];
// Block on the async call using our global runtime
let response = get_runtime().block_on(async {
let mut stream = match agent.reply(&messages, None, None).await {
Ok(stream) => stream,
Err(e) => return format!("Error getting reply from agent: {}", e),
};
let mut full_response = String::new();
while let Some(message_result) = stream.next().await {
match message_result {
Ok(AgentEvent::Message(message)) => {
// Get text or serialize to JSON
// Note: Message doesn't have as_text method, we'll serialize to JSON
if let Ok(json) = serde_json::to_string(&message) {
full_response.push_str(&json);
}
}
Ok(AgentEvent::McpNotification(_)) => {
// TODO: Handle MCP notifications.
}
Ok(AgentEvent::ModelChange { .. }) => {
// Model change events are informational, just continue
}
Err(e) => {
full_response.push_str(&format!("\nError in message stream: {}", e));
}
}
}
full_response
});
string_to_c_char(&response)
}
// Tool schema creation will be implemented in a future commit
/// Free a string allocated by goose FFI functions
///
/// This function frees memory allocated for strings returned by goose FFI functions.
///
/// # Parameters
///
/// - s: String to free
///
/// # Safety
///
/// The string must have been allocated by a goose FFI function, or be NULL.
/// The string must not be used after calling this function.
#[no_mangle]
pub unsafe extern "C" fn goose_free_string(s: *mut c_char) {
if !s.is_null() {
let _ = CString::from_raw(s);
}
}
// Helper function to convert a Rust string to a C char pointer
fn string_to_c_char(s: &str) -> *mut c_char {
match CString::new(s) {
Ok(c_string) => c_string.into_raw(),
Err(_) => ptr::null_mut(),
}
}
-73
View File
@@ -1,73 +0,0 @@
[package]
name = "goose-llm"
edition.workspace = true
version.workspace = true
authors.workspace = true
license.workspace = true
repository.workspace = true
description.workspace = true
[lints]
workspace = true
[lib]
crate-type = ["lib", "cdylib"]
name = "goose_llm"
[dependencies]
goose = { path = "../goose" }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
anyhow = "1.0"
thiserror = "1.0"
minijinja = "2.8.0"
include_dir = "0.7.4"
once_cell = "1.20.2"
chrono = { version = "0.4.38", features = ["serde"] }
reqwest = { version = "0.12.9", features = [
"rustls-tls-native-roots",
"json",
"cookies",
"gzip",
"brotli",
"deflate",
"zstd",
"charset",
"http2",
"stream"
], default-features = false }
async-trait = "0.1"
url = "2.5"
base64 = "0.21"
regex = "1.11.1"
tracing = "0.1"
smallvec = { version = "1.13", features = ["serde"] }
indoc = "1.0"
# https://github.com/mozilla/uniffi-rs/blob/c7f6caa3d1bf20f934346cefd8e82b5093f0dc6f/fixtures/futures/Cargo.toml#L22
uniffi = { version = "0.29", features = ["tokio", "cli", "scaffolding-ffi-buffer-fns"] }
tokio = { version = "1.43", features = ["time", "sync"] }
[dev-dependencies]
criterion = "0.5"
tempfile = "3.15.0"
dotenvy = "0.15.7"
lazy_static = "1.5"
ctor = "0.2.7"
tokio = { version = "1.43", features = ["full"] }
[[bin]]
# https://mozilla.github.io/uniffi-rs/latest/tutorial/foreign_language_bindings.html
name = "uniffi-bindgen"
path = "uniffi-bindgen.rs"
[[example]]
name = "simple"
path = "examples/simple.rs"
[[example]]
name = "image"
path = "examples/image.rs"
[[example]]
name = "prompt_override"
path = "examples/prompt_override.rs"
-80
View File
@@ -1,80 +0,0 @@
## goose-llm
This crate is meant to be used for foreign function interface (FFI). It's meant to be
stateless and contain logic related to providers and prompts:
- chat completion with model providers
- detecting read-only tools for smart approval
- methods for summarization / truncation
Run:
```
cargo run -p goose-llm --example simple
```
## Kotlin bindings
Structure:
```
.
└── crates
└── goose-llm/...
└── target
└── debug/libgoose_llm.dylib
├── bindings
│ └── kotlin
│ ├── example
│ │ └── Usage.kt ← your demo app
│ └── uniffi
│ └── goose_llm
│ └── goose_llm.kt ← auto-generated bindings
```
#### Kotlin -> Rust: run example
The following `just` command creates kotlin bindings, then compiles and runs an example.
```bash
just kotlin-example
```
You will have to download jars in `bindings/kotlin/libs` directory (only the first time):
```bash
pushd bindings/kotlin/libs/
curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlin/kotlin-stdlib/1.9.0/kotlin-stdlib-1.9.0.jar
curl -O https://repo1.maven.org/maven2/org/jetbrains/kotlinx/kotlinx-coroutines-core-jvm/1.7.3/kotlinx-coroutines-core-jvm-1.7.3.jar
curl -O https://repo1.maven.org/maven2/net/java/dev/jna/jna/5.13.0/jna-5.13.0.jar
popd
```
To just create the Kotlin bindings (for MacOS):
```bash
# run from project root directory
cargo build -p goose-llm
cargo run --features=uniffi/cli --bin uniffi-bindgen generate --library ./target/debug/libgoose_llm.dylib --language kotlin --out-dir bindings/kotlin
```
Creating `libgoose_llm.so` for Linux distribution:
Use `cross` to build for the specific target and then create the bindings:
```
# x86-64 GNU/Linux (kGoose uses this)
rustup target add x86_64-unknown-linux-gnu
cross build --release --target x86_64-unknown-linux-gnu -p goose-llm
# The goose_llm.kt bindings produced should be the same whether we use 'libgoose_llm.dylib' or 'libgoose_llm.so'
cross run --features=uniffi/cli --bin uniffi-bindgen generate --library ./target/x86_64-unknown-linux-gnu/release/libgoose_llm.so --language kotlin --out-dir bindings/kotlin
```
#### Python -> Rust: generate bindings, run example
```bash
cargo run --features=uniffi/cli --bin uniffi-bindgen generate --library ./target/debug/libgoose_llm.dylib --language python --out-dir bindings/python
DYLD_LIBRARY_PATH=./target/debug python bindings/python/usage.py
```
-53
View File
@@ -1,53 +0,0 @@
use anyhow::Result;
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use goose_llm::{
completion,
message::MessageContent,
types::completion::{CompletionRequest, CompletionResponse},
Message, ModelConfig,
};
use serde_json::json;
use std::{fs, vec};
#[tokio::main]
async fn main() -> Result<()> {
let provider = "databricks";
let provider_config = json!({
"host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
"token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
});
let model_name = "goose-claude-4-sonnet"; // "gpt-4o";
let model_config = ModelConfig::new(model_name.to_string());
let system_preamble = "You are a helpful assistant.";
// Read and encode test image
let image_data = fs::read("examples/test_assets/test_image.png")?;
let base64_image = BASE64.encode(image_data);
let user_msg = Message::user()
.with_text("What do you see in this image?")
.with_content(MessageContent::image(base64_image, "image/png"));
let messages = vec![user_msg];
let completion_response: CompletionResponse = completion(
CompletionRequest::new(
provider.to_string(),
provider_config.clone(),
model_config.clone(),
Some(system_preamble.to_string()),
None,
messages,
vec![],
)
.with_request_id("test-image-1".to_string()),
)
.await?;
// Print the response
println!("\nCompletion Response:");
println!("{}", serde_json::to_string_pretty(&completion_response)?);
Ok(())
}
@@ -1,48 +0,0 @@
use std::vec;
use anyhow::Result;
use goose_llm::{
completion,
types::completion::{CompletionRequest, CompletionResponse},
Message, ModelConfig,
};
use serde_json::json;
#[tokio::main]
async fn main() -> Result<()> {
let provider = "databricks";
let provider_config = json!({
"host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
"token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
});
// let model_name = "goose-gpt-4-1"; // parallel tool calls
let model_name = "claude-3-5-haiku";
let model_config = ModelConfig::new(model_name.to_string());
let system_prompt_override = "You are a helpful assistant. Talk in the style of pirates.";
for text in ["How was your day?"] {
println!("\n---------------\n");
println!("User Input: {text}");
let messages = vec![
Message::user().with_text("Hi there!"),
Message::assistant().with_text("How can I help?"),
Message::user().with_text(text),
];
let completion_response: CompletionResponse = completion(CompletionRequest::new(
provider.to_string(),
provider_config.clone(),
model_config.clone(),
None,
Some(system_prompt_override.to_string()),
messages.clone(),
vec![],
))
.await?;
// Print the response
println!("\nCompletion Response:");
println!("{}", serde_json::to_string_pretty(&completion_response)?);
}
Ok(())
}
-124
View File
@@ -1,124 +0,0 @@
use std::vec;
use anyhow::Result;
use goose_llm::{
completion,
extractors::generate_tooltip,
types::completion::{
CompletionRequest, CompletionResponse, ExtensionConfig, ToolApprovalMode, ToolConfig,
},
Message, ModelConfig,
};
use serde_json::json;
#[tokio::main]
async fn main() -> Result<()> {
let provider = "databricks";
let provider_config = json!({
"host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
"token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
});
// let model_name = "goose-gpt-4-1"; // parallel tool calls
let model_name = "claude-3-5-haiku";
let model_config = ModelConfig::new(model_name.to_string());
let calculator_tool = ToolConfig::new(
"calculator",
"Perform basic arithmetic operations",
json!({
"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",
}
}
}),
ToolApprovalMode::Auto,
);
let bash_tool = ToolConfig::new(
"bash_shell",
"Run a shell command",
json!({
"type": "object",
"required": ["command"],
"properties": {
"command": {
"type": "string",
"description": "The shell command to execute"
}
}
}),
ToolApprovalMode::Manual,
);
let list_dir_tool = ToolConfig::new(
"list_directory",
"List files in a directory",
json!({
"type": "object",
"required": ["path"],
"properties": {
"path": {
"type": "string",
"description": "The directory path to list files from"
}
}
}),
ToolApprovalMode::Auto,
);
let extensions = vec![
ExtensionConfig::new(
"calculator_extension".to_string(),
Some("This extension provides a calculator tool.".to_string()),
vec![calculator_tool],
),
ExtensionConfig::new(
"bash_extension".to_string(),
Some("This extension provides a bash shell tool.".to_string()),
vec![bash_tool, list_dir_tool],
),
];
let system_preamble = "You are a helpful assistant.";
for text in [
"Add 10037 + 23123 using calculator and also run 'date -u' using bash",
"List all files in the current directory",
] {
println!("\n---------------\n");
println!("User Input: {text}");
let messages = vec![
Message::user().with_text("Hi there!"),
Message::assistant().with_text("How can I help?"),
Message::user().with_text(text),
];
let completion_response: CompletionResponse = completion(CompletionRequest::new(
provider.to_string(),
provider_config.clone(),
model_config.clone(),
Some(system_preamble.to_string()),
None,
messages.clone(),
extensions.clone(),
))
.await?;
// Print the response
println!("\nCompletion Response:");
println!("{}", serde_json::to_string_pretty(&completion_response)?);
let tooltip = generate_tooltip(provider, provider_config.clone(), &messages, None).await?;
println!("\nTooltip: {}", tooltip);
}
Ok(())
}
Binary file not shown.

Before

Width:  |  Height:  |  Size: 4.2 KiB

-168
View File
@@ -1,168 +0,0 @@
use std::{collections::HashMap, time::Instant};
use anyhow::Result;
use chrono::Utc;
use serde_json::Value;
use crate::{
message::{Message, MessageContent},
prompt_template,
providers::create,
types::{
completion::{
CompletionError, CompletionRequest, CompletionResponse, ExtensionConfig,
RuntimeMetrics, ToolApprovalMode, ToolConfig,
},
core::ToolCall,
},
};
#[uniffi::export]
pub fn print_messages(messages: Vec<Message>) {
for msg in messages {
println!("[{:?} @ {}] {:?}", msg.role, msg.created, msg.content);
}
}
/// Public API for the Goose LLM completion function
#[uniffi::export(async_runtime = "tokio")]
pub async fn completion(req: CompletionRequest) -> Result<CompletionResponse, CompletionError> {
let start_total = Instant::now();
let provider = create(
&req.provider_name,
req.provider_config.clone(),
req.model_config.clone(),
)
.map_err(|_| CompletionError::UnknownProvider(req.provider_name.to_string()))?;
let system_prompt = construct_system_prompt(
&req.system_preamble,
&req.system_prompt_override,
&req.extensions,
)?;
let tools = collect_prefixed_tools(&req.extensions);
// Call the LLM provider
let start_provider = Instant::now();
let mut response = provider
.complete(
&system_prompt,
&req.messages,
&tools,
req.request_id.as_deref(),
)
.await?;
let provider_elapsed_sec = start_provider.elapsed().as_secs_f32();
let usage_tokens = response.usage.total_tokens;
let tool_configs = collect_prefixed_tool_configs(&req.extensions);
update_needs_approval_for_tool_calls(&mut response.message, &tool_configs)?;
Ok(CompletionResponse::new(
response.message,
response.model,
response.usage,
calculate_runtime_metrics(start_total, provider_elapsed_sec, usage_tokens),
))
}
/// Render the global `system.md` template with the provided context.
fn construct_system_prompt(
preamble: &Option<String>,
prompt_override: &Option<String>,
extensions: &[ExtensionConfig],
) -> Result<String, CompletionError> {
// If both system_preamble and system_prompt_override are provided, then prompt_override takes precedence
// and we don't render the template using preamble and extensions. Just return the prompt_override as is.
if prompt_override.is_some() {
return Ok(prompt_override.clone().unwrap());
}
let system_preamble = {
if let Some(p) = preamble {
p
} else {
"You are a helpful assistant."
}
};
let mut context: HashMap<&str, Value> = HashMap::new();
context.insert("system_preamble", Value::String(system_preamble.to_owned()));
context.insert("extensions", serde_json::to_value(extensions)?);
context.insert(
"current_date",
Value::String(Utc::now().format("%Y-%m-%d").to_string()),
);
Ok(prompt_template::render_global_file("system.md", &context)?)
}
/// Determine if a tool call requires manual approval.
fn determine_needs_approval(config: &ToolConfig, _call: &ToolCall) -> bool {
match config.approval_mode {
ToolApprovalMode::Auto => false,
ToolApprovalMode::Manual => true,
ToolApprovalMode::Smart => {
// TODO: Implement smart approval logic later
true
}
}
}
/// Set `needs_approval` on every tool call in the message.
/// Returns a `ToolNotFound` error if the corresponding `ToolConfig` is missing.
pub fn update_needs_approval_for_tool_calls(
message: &mut Message,
tool_configs: &HashMap<String, ToolConfig>,
) -> Result<(), CompletionError> {
for content in &mut message.content.iter_mut() {
if let MessageContent::ToolReq(req) = content {
if let Ok(call) = &mut req.tool_call.0 {
// Provide a clear error message when the tool config is missing
let config = tool_configs.get(&call.name).ok_or_else(|| {
CompletionError::ToolNotFound(format!(
"could not find tool config for '{}'",
call.name
))
})?;
let needs_approval = determine_needs_approval(config, call);
call.set_needs_approval(needs_approval);
}
}
}
Ok(())
}
/// Collect all `Tool` instances from the extensions.
fn collect_prefixed_tools(extensions: &[ExtensionConfig]) -> Vec<crate::types::core::Tool> {
extensions
.iter()
.flat_map(|ext| ext.get_prefixed_tools())
.collect()
}
/// Collect all `ToolConfig` entries from the extensions into a map.
fn collect_prefixed_tool_configs(extensions: &[ExtensionConfig]) -> HashMap<String, ToolConfig> {
extensions
.iter()
.flat_map(|ext| ext.get_prefixed_tool_configs())
.collect()
}
/// Compute runtime metrics for the request.
fn calculate_runtime_metrics(
total_start: Instant,
provider_elapsed_sec: f32,
token_count: Option<i32>,
) -> RuntimeMetrics {
let total_ms = total_start.elapsed().as_secs_f32();
let tokens_per_sec = token_count.and_then(|toks| {
if provider_elapsed_sec > 0.0 {
Some(toks as f64 / (provider_elapsed_sec as f64))
} else {
None
}
});
RuntimeMetrics::new(total_ms, provider_elapsed_sec, tokens_per_sec)
}
-5
View File
@@ -1,5 +0,0 @@
mod session_name;
mod tooltip;
pub use session_name::generate_session_name;
pub use tooltip::generate_tooltip;
@@ -1,110 +0,0 @@
use crate::generate_structured_outputs;
use crate::providers::errors::ProviderError;
use crate::types::core::Role;
use crate::{message::Message, types::json_value_ffi::JsonValueFfi};
use anyhow::Result;
use goose::utils::safe_truncate;
use indoc::indoc;
use serde_json::{json, Value};
const SESSION_NAME_EXAMPLES: &[&str] = &[
"Research Synthesis",
"Sentiment Analysis",
"Performance Report",
"Feedback Collector",
"Accessibility Check",
"Design Reminder",
"Project Reminder",
"Launch Checklist",
"Metrics Monitor",
"Incident Response",
"Deploy Cabinet App",
"Design Reminder Alert",
"Generate Monthly Expense Report",
"Automate Incident Response Workflow",
"Analyze Brand Sentiment Trends",
"Monitor Device Health Issues",
"Collect UI Feedback Summary",
"Schedule Project Deadline Reminders",
];
fn build_system_prompt() -> String {
let examples = SESSION_NAME_EXAMPLES
.iter()
.map(|e| format!("- {}", e))
.collect::<Vec<_>>()
.join("\n");
indoc! {r#"
You are an assistant that crafts a concise session title.
Given the first couple user messages in the conversation so far,
reply with only a short name (up to 4 words) that best describes
this session's goal.
Examples:
"#}
.to_string()
+ &examples
}
/// Generates a short (≤4 words) session name
#[uniffi::export(async_runtime = "tokio", default(request_id = None))]
pub async fn generate_session_name(
provider_name: &str,
provider_config: JsonValueFfi,
messages: &[Message],
request_id: Option<String>,
) -> Result<String, ProviderError> {
// Collect up to the first 3 user messages (truncated to 300 chars each)
let context: Vec<String> = messages
.iter()
.filter(|m| m.role == Role::User)
.take(3)
.map(|m| {
let text = m.content.concat_text_str();
safe_truncate(&text, 300)
})
.collect();
if context.is_empty() {
return Err(ProviderError::ExecutionError(
"No user messages found to generate a session name.".to_string(),
));
}
let system_prompt = build_system_prompt();
let user_msg_text = format!("Here are the user messages:\n{}", context.join("\n"));
// Use `extract` with a simple string schema
let schema = json!({
"type": "object",
"properties": {
"name": { "type": "string" }
},
"required": ["name"],
"additionalProperties": false
});
let resp = generate_structured_outputs(
provider_name,
provider_config,
&system_prompt,
&[Message::user().with_text(&user_msg_text)],
schema,
request_id,
)
.await?;
let obj = resp
.data
.as_object()
.ok_or_else(|| ProviderError::ResponseParseError("Expected object".into()))?;
let name = obj
.get("name")
.and_then(Value::as_str)
.ok_or_else(|| ProviderError::ResponseParseError("Missing or non-string name".into()))?
.to_string();
Ok(name)
}
-171
View File
@@ -1,171 +0,0 @@
use crate::generate_structured_outputs;
use crate::message::{Message, MessageContent};
use crate::providers::errors::ProviderError;
use crate::types::core::{Content, Role};
use crate::types::json_value_ffi::JsonValueFfi;
use anyhow::Result;
use indoc::indoc;
use serde_json::{json, Value};
const TOOLTIP_EXAMPLES: &[&str] = &[
"analyzing KPIs",
"detecting anomalies",
"building artifacts in Buildkite",
"categorizing issues",
"checking dependencies",
"collecting feedback",
"deploying changes in AWS",
"drafting report in Google Docs",
"extracting action items",
"generating insights",
"logging issues",
"monitoring tickets in Zendesk",
"notifying design team",
"running integration tests",
"scanning threads in Figma",
"sending reminders in Gmail",
"sending surveys",
"sharing with stakeholders",
"summarizing findings",
"transcribing meeting",
"tracking resolution",
"updating status in Linear",
];
fn build_system_prompt() -> String {
let examples = TOOLTIP_EXAMPLES
.iter()
.map(|e| format!("- {}", e))
.collect::<Vec<_>>()
.join("\n");
indoc! {r#"
You are an assistant that summarizes the recent conversation into a tooltip.
Given the last two messages, reply with only a short tooltip (up to 4 words)
describing what is happening now.
Examples:
"#}
.to_string()
+ &examples
}
/// Generates a tooltip summarizing the last two messages in the session,
/// including any tool calls or results.
#[uniffi::export(async_runtime = "tokio", default(request_id = None))]
pub async fn generate_tooltip(
provider_name: &str,
provider_config: JsonValueFfi,
messages: &[Message],
request_id: Option<String>,
) -> Result<String, ProviderError> {
// Need at least two messages to generate a tooltip
if messages.len() < 2 {
return Err(ProviderError::ExecutionError(
"Need at least two messages to generate a tooltip".to_string(),
));
}
// Helper to render a single message's content
fn render_message(m: &Message) -> String {
let mut parts = Vec::new();
for content in m.content.iter() {
match content {
MessageContent::Text(text_block) => {
let txt = text_block.text.trim();
if !txt.is_empty() {
parts.push(txt.to_string());
}
}
MessageContent::ToolReq(req) => {
if let Ok(tool_call) = &req.tool_call.0 {
parts.push(format!(
"called tool '{}' with args {}",
tool_call.name, tool_call.arguments
));
} else if let Err(e) = &req.tool_call.0 {
parts.push(format!("tool request error: {}", e));
}
}
MessageContent::ToolResp(resp) => match &resp.tool_result.0 {
Ok(contents) => {
let results: Vec<String> = contents
.iter()
.map(|c| match c {
Content::Text(t) => t.text.clone(),
Content::Image(_) => "[image]".to_string(),
})
.collect();
parts.push(format!("tool responded with: {}", results.join(" ")));
}
Err(e) => {
parts.push(format!("tool error: {}", e));
}
},
_ => {} // ignore other variants
}
}
let role = match m.role {
Role::User => "User",
Role::Assistant => "Assistant",
};
format!("{}: {}", role, parts.join("; "))
}
// Take the last two messages (in correct chronological order)
let rendered: Vec<String> = messages
.iter()
.rev()
.take(2)
.map(render_message)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
let system_prompt = build_system_prompt();
let user_msg_text = format!(
"Here are the last two messages:\n{}\n\nTooltip:",
rendered.join("\n")
);
// Schema wrapping our tooltip string
let schema = json!({
"type": "object",
"properties": {
"tooltip": { "type": "string" }
},
"required": ["tooltip"],
"additionalProperties": false
});
// Get the structured outputs
let resp = generate_structured_outputs(
provider_name,
provider_config,
&system_prompt,
&[Message::user().with_text(&user_msg_text)],
schema,
request_id,
)
.await?;
// Pull out the tooltip field
let obj = resp
.data
.as_object()
.ok_or_else(|| ProviderError::ResponseParseError("Expected JSON object".into()))?;
let tooltip = obj
.get("tooltip")
.and_then(Value::as_str)
.ok_or_else(|| {
ProviderError::ResponseParseError("Missing or non-string `tooltip` field".into())
})?
.to_string();
Ok(tooltip)
}
-15
View File
@@ -1,15 +0,0 @@
uniffi::setup_scaffolding!();
mod completion;
pub mod extractors;
pub mod message;
mod model;
mod prompt_template;
pub mod providers;
mod structured_outputs;
pub mod types;
pub use completion::completion;
pub use message::Message;
pub use model::ModelConfig;
pub use structured_outputs::generate_structured_outputs;
-186
View File
@@ -1,186 +0,0 @@
use std::{iter::FromIterator, ops::Deref};
use crate::message::MessageContent;
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
/// Holds the heterogeneous fragments that make up one chat message.
///
/// * Up to two items are stored inline on the stack.
/// * Falls back to a heap allocation only when necessary.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(transparent)]
pub struct Contents(SmallVec<[MessageContent; 2]>);
impl Contents {
/*----------------------------------------------------------
* 1-line ergonomic helpers
*---------------------------------------------------------*/
pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, MessageContent> {
self.0.iter_mut()
}
pub fn push(&mut self, item: impl Into<MessageContent>) {
self.0.push(item.into());
}
pub fn texts(&self) -> impl Iterator<Item = &str> {
self.0.iter().filter_map(|c| c.as_text())
}
pub fn concat_text_str(&self) -> String {
self.texts().collect::<Vec<_>>().join("\n")
}
/// Returns `true` if *any* item satisfies the predicate.
pub fn any_is<P>(&self, pred: P) -> bool
where
P: FnMut(&MessageContent) -> bool,
{
self.iter().any(pred)
}
/// Returns `true` if *every* item satisfies the predicate.
pub fn all_are<P>(&self, pred: P) -> bool
where
P: FnMut(&MessageContent) -> bool,
{
self.iter().all(pred)
}
}
impl From<Vec<MessageContent>> for Contents {
fn from(v: Vec<MessageContent>) -> Self {
Contents(SmallVec::from_vec(v))
}
}
impl FromIterator<MessageContent> for Contents {
fn from_iter<I: IntoIterator<Item = MessageContent>>(iter: I) -> Self {
Contents(SmallVec::from_iter(iter))
}
}
/*--------------------------------------------------------------
* Allow &message.content to behave like a slice of fragments.
*-------------------------------------------------------------*/
impl Deref for Contents {
type Target = [MessageContent];
fn deref(&self) -> &Self::Target {
&self.0
}
}
// — Register the contents type with UniFFI, converting to/from Vec<MessageContent> —
// We need to do this because UniFFIs FFI layer supports only primitive buffers (here Vec<u8>),
uniffi::custom_type!(Contents, Vec<MessageContent>, {
lower: |contents: &Contents| {
contents.0.to_vec()
},
try_lift: |contents: Vec<MessageContent>| {
Ok(Contents::from(contents))
},
});
#[cfg(test)]
mod tests {
use super::*;
use crate::types::core::{Content, TextContent, ToolCall, ToolError};
use serde_json::json;
// ------------------------------------------------------------
// Helpers
// ------------------------------------------------------------
fn make_tool_req_ok(id: &str) -> MessageContent {
let call = ToolCall::new("echo", json!({"text": "hi"}));
MessageContent::tool_request(id, Ok(call).into())
}
fn make_tool_resp_ok(id: &str) -> MessageContent {
let body = vec![Content::Text(TextContent {
text: "done".into(),
})];
MessageContent::tool_response(id, Ok(body).into())
}
fn make_tool_req_err(id: &str) -> MessageContent {
let err = ToolError::NotFound(format!(
"The provided function name '{}' had invalid characters",
"bad$name"
));
MessageContent::tool_request(id, Err(err).into())
}
fn make_tool_resp_err(id: &str) -> MessageContent {
let err = ToolError::InvalidParameters("Could not interpret tool use parameters".into());
MessageContent::tool_response(id, Err(err).into())
}
// ------------------------------------------------------------
// Round-trip: success
// ------------------------------------------------------------
#[test]
fn contents_roundtrip_ok() {
let items: Contents = vec![make_tool_req_ok("req-1"), make_tool_resp_ok("resp-1")].into();
// ---- serialise
let json_str = serde_json::to_string(&items).expect("serialise OK");
println!("JSON: {:?}", json_str);
assert!(
json_str.contains(r#""type":"toolReq""#)
&& json_str.contains(r#""type":"toolResp""#)
&& json_str.contains(r#""status":"success""#),
"JSON should contain both variants and success-status"
);
// ---- deserialise
let parsed: Contents = serde_json::from_str(&json_str).expect("deserialise OK");
assert_eq!(parsed, items, "full round-trip equality");
}
// ------------------------------------------------------------
// Round-trip: error (all variants collapse to ExecutionError)
// ------------------------------------------------------------
#[test]
fn contents_roundtrip_err() {
let original_items: Contents =
vec![make_tool_req_err("req-e"), make_tool_resp_err("resp-e")].into();
// ---- serialise
let json_str = serde_json::to_string(&original_items).expect("serialise OK");
println!("JSON: {:?}", json_str);
assert!(json_str.contains(r#""status":"error""#));
// ---- deserialise
let parsed: Contents = serde_json::from_str(&json_str).expect("deserialise OK");
// ─── validate structure ───────────────────────────────────
assert_eq!(parsed.len(), 2);
// ToolReq error
match &parsed[0] {
MessageContent::ToolReq(req) => match &*req.tool_call {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("invalid characters"))
}
other => panic!("expected ExecutionError, got {:?}", other),
},
other => panic!("expected ToolReq, got {:?}", other),
}
// ToolResp error
match &parsed[1] {
MessageContent::ToolResp(resp) => match &*resp.tool_result {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("interpret tool use parameters"))
}
other => panic!("expected ExecutionError, got {:?}", other),
},
other => panic!("expected ToolResp, got {:?}", other),
}
}
}
@@ -1,465 +0,0 @@
use serde::{Deserialize, Serialize};
use serde_json::{self, Deserializer, Serializer};
use crate::message::tool_result_serde;
use crate::types::core::{Content, ImageContent, TextContent, ToolCall, ToolResult};
// — Newtype wrappers (local structs) so we satisfy Rusts orphan rules —
// We need these because we cant implement UniFFIs FfiConverter directly on a type alias.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolRequestToolCall(#[serde(with = "tool_result_serde")] pub ToolResult<ToolCall>);
impl ToolRequestToolCall {
pub fn as_result(&self) -> &ToolResult<ToolCall> {
&self.0
}
}
impl std::ops::Deref for ToolRequestToolCall {
type Target = ToolResult<ToolCall>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<Result<ToolCall, crate::types::core::ToolError>> for ToolRequestToolCall {
fn from(res: Result<ToolCall, crate::types::core::ToolError>) -> Self {
ToolRequestToolCall(res)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolResponseToolResult(
#[serde(with = "tool_result_serde")] pub ToolResult<Vec<Content>>,
);
impl ToolResponseToolResult {
pub fn as_result(&self) -> &ToolResult<Vec<Content>> {
&self.0
}
}
impl std::ops::Deref for ToolResponseToolResult {
type Target = ToolResult<Vec<Content>>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<Result<Vec<Content>, crate::types::core::ToolError>> for ToolResponseToolResult {
fn from(res: Result<Vec<Content>, crate::types::core::ToolError>) -> Self {
ToolResponseToolResult(res)
}
}
// — Register the newtypes with UniFFI, converting via JSON strings —
// UniFFIs FFI layer supports only primitive buffers (here String), so we JSON-serialize
// through our `tool_result_serde` to preserve the same success/error schema on both sides.
// see https://github.com/mozilla/uniffi-rs/issues/2533
uniffi::custom_type!(ToolRequestToolCall, String, {
lower: |wrapper: &ToolRequestToolCall| {
let mut buf = Vec::new();
{
let mut ser = Serializer::new(&mut buf);
// note the borrow on wrapper.0
tool_result_serde::serialize(&wrapper.0, &mut ser)
.expect("ToolRequestToolCall serialization failed");
}
String::from_utf8(buf).expect("ToolRequestToolCall produced invalid UTF-8")
},
try_lift: |s: String| {
let mut de = Deserializer::from_str(&s);
let result = tool_result_serde::deserialize(&mut de)
.map_err(anyhow::Error::new)?;
Ok(ToolRequestToolCall(result))
},
});
uniffi::custom_type!(ToolResponseToolResult, String, {
lower: |wrapper: &ToolResponseToolResult| {
let mut buf = Vec::new();
{
let mut ser = Serializer::new(&mut buf);
// note the borrow on wrapper.0
tool_result_serde::serialize(&wrapper.0, &mut ser)
.expect("ToolResponseToolResult serialization failed");
}
String::from_utf8(buf).expect("ToolResponseToolResult produced invalid UTF-8")
},
try_lift: |s: String| {
let mut de = Deserializer::from_str(&s);
let result = tool_result_serde::deserialize(&mut de)
.map_err(anyhow::Error::new)?;
Ok(ToolResponseToolResult(result))
},
});
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
#[serde(rename_all = "camelCase")]
pub struct ToolRequest {
pub id: String,
pub tool_call: ToolRequestToolCall,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
#[serde(rename_all = "camelCase")]
pub struct ToolResponse {
pub id: String,
pub tool_result: ToolResponseToolResult,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
pub struct ThinkingContent {
pub thinking: String,
pub signature: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
pub struct RedactedThinkingContent {
pub data: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Enum)]
/// Content passed inside a message, which can be both simple content and tool content
#[serde(tag = "type", rename_all = "camelCase")]
pub enum MessageContent {
Text(TextContent),
Image(ImageContent),
ToolReq(ToolRequest),
ToolResp(ToolResponse),
Thinking(ThinkingContent),
RedactedThinking(RedactedThinkingContent),
}
impl MessageContent {
pub fn text<S: Into<String>>(text: S) -> Self {
MessageContent::Text(TextContent { text: text.into() })
}
pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
MessageContent::Image(ImageContent {
data: data.into(),
mime_type: mime_type.into(),
})
}
pub fn tool_request<S: Into<String>>(id: S, tool_call: ToolRequestToolCall) -> Self {
MessageContent::ToolReq(ToolRequest {
id: id.into(),
tool_call,
})
}
pub fn tool_response<S: Into<String>>(id: S, tool_result: ToolResponseToolResult) -> Self {
MessageContent::ToolResp(ToolResponse {
id: id.into(),
tool_result,
})
}
pub fn thinking<S1: Into<String>, S2: Into<String>>(thinking: S1, signature: S2) -> Self {
MessageContent::Thinking(ThinkingContent {
thinking: thinking.into(),
signature: signature.into(),
})
}
pub fn redacted_thinking<S: Into<String>>(data: S) -> Self {
MessageContent::RedactedThinking(RedactedThinkingContent { data: data.into() })
}
pub fn as_tool_request(&self) -> Option<&ToolRequest> {
if let MessageContent::ToolReq(ref tool_request) = self {
Some(tool_request)
} else {
None
}
}
pub fn as_tool_response(&self) -> Option<&ToolResponse> {
if let MessageContent::ToolResp(ref tool_response) = self {
Some(tool_response)
} else {
None
}
}
pub fn as_tool_response_text(&self) -> Option<String> {
if let Some(tool_response) = self.as_tool_response() {
if let Ok(contents) = &tool_response.tool_result.0 {
let texts: Vec<String> = contents
.iter()
.filter_map(|content| content.as_text().map(String::from))
.collect();
if !texts.is_empty() {
return Some(texts.join("\n"));
}
}
}
None
}
pub fn as_tool_request_id(&self) -> Option<&str> {
if let Self::ToolReq(r) = self {
Some(&r.id)
} else {
None
}
}
pub fn as_tool_response_id(&self) -> Option<&str> {
if let Self::ToolResp(r) = self {
Some(&r.id)
} else {
None
}
}
/// Get the text content if this is a TextContent variant
pub fn as_text(&self) -> Option<&str> {
match self {
MessageContent::Text(text) => Some(&text.text),
_ => None,
}
}
/// Get the thinking content if this is a ThinkingContent variant
pub fn as_thinking(&self) -> Option<&ThinkingContent> {
match self {
MessageContent::Thinking(thinking) => Some(thinking),
_ => None,
}
}
/// Get the redacted thinking content if this is a RedactedThinkingContent variant
pub fn as_redacted_thinking(&self) -> Option<&RedactedThinkingContent> {
match self {
MessageContent::RedactedThinking(redacted) => Some(redacted),
_ => None,
}
}
pub fn is_text(&self) -> bool {
matches!(self, Self::Text(_))
}
pub fn is_image(&self) -> bool {
matches!(self, Self::Image(_))
}
pub fn is_tool_request(&self) -> bool {
matches!(self, Self::ToolReq(_))
}
pub fn is_tool_response(&self) -> bool {
matches!(self, Self::ToolResp(_))
}
}
impl From<Content> for MessageContent {
fn from(content: Content) -> Self {
match content {
Content::Text(text) => MessageContent::Text(text),
Content::Image(image) => MessageContent::Image(image),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::core::{ToolCall, ToolError};
use crate::UniFfiTag;
use serde_json::json;
use uniffi::{FfiConverter, RustBuffer};
// ---------- ToolRequestToolCall ----------------------------------------------------------
#[test]
fn tool_request_tool_call_roundtrip_ok() {
// Build a valid ToolCall
let call = ToolCall::new("my_function", json!({"a": 1, "b": 2}));
// Wrap it in the new-type
let wrapper = ToolRequestToolCall::from(Ok(call.clone()));
// Serialize → JSON
let json_str = serde_json::to_string(&wrapper).expect("serialize OK");
assert!(
json_str.contains(r#""status":"success""#),
"must mark success"
);
// Deserialize ← JSON
let parsed: ToolRequestToolCall = serde_json::from_str(&json_str).expect("deserialize OK");
// Round-trip equality
assert_eq!(*parsed, Ok(call));
}
#[test]
fn tool_request_tool_call_roundtrip_err() {
// Typical failure variant that could come from `is_valid_function_name`
let err = ToolError::NotFound(
"The provided function name 'bad$name' had invalid characters".into(),
);
let wrapper = ToolRequestToolCall::from(Err(err.clone()));
let json_str = serde_json::to_string(&wrapper).expect("serialize OK");
assert!(
json_str.contains(r#""status":"error""#) && json_str.contains("invalid characters"),
"must mark error and carry message"
);
let parsed: ToolRequestToolCall = serde_json::from_str(&json_str).expect("deserialize OK");
match &*parsed {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("invalid characters"))
}
other => panic!("expected ExecutionError, got {:?}", other),
}
}
// ---------- ToolResponseToolResult -------------------------------------------------------
#[test]
fn tool_response_tool_result_roundtrip_ok() {
// Minimal content vector (one text item)
let content_vec = vec![Content::Text(TextContent {
text: "hello".into(),
})];
let wrapper = ToolResponseToolResult::from(Ok(content_vec.clone()));
let json_str = serde_json::to_string(&wrapper).expect("serialize OK");
assert!(json_str.contains(r#""status":"success""#));
let parsed: ToolResponseToolResult =
serde_json::from_str(&json_str).expect("deserialize OK");
assert_eq!(*parsed, Ok(content_vec));
}
#[test]
fn tool_response_tool_result_roundtrip_err() {
let err = ToolError::InvalidParameters("Could not interpret tool use parameters".into());
let wrapper = ToolResponseToolResult::from(Err(err.clone()));
let json_str = serde_json::to_string(&wrapper).expect("serialize OK");
assert!(json_str.contains(r#""status":"error""#));
let parsed: ToolResponseToolResult =
serde_json::from_str(&json_str).expect("deserialize OK");
match &*parsed {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("interpret tool use"))
}
other => panic!("expected ExecutionError, got {:?}", other),
}
}
// ---------- FFI (lower / lift) round-trips ----------------------------------------------
// https://mozilla.github.io/uniffi-rs/latest/internals/lifting_and_lowering.html
#[test]
fn ffi_roundtrip_tool_request_ok_and_err() {
// ---------- status: success ----------
let ok_call = ToolCall::new("echo", json!({"text": "hi"}));
let ok_wrapper = ToolRequestToolCall::from(Ok(ok_call.clone()));
// First lower → inspect JSON
let buf1: RustBuffer =
<ToolRequestToolCall as FfiConverter<UniFfiTag>>::lower(ok_wrapper.clone());
let json_ok: String =
<String as FfiConverter<UniFfiTag>>::try_lift(buf1).expect("lift String OK");
println!("ToolReq - Lowered JSON (status: success): {:?}", json_ok);
assert!(json_ok.contains(r#""status":"success""#));
// Second lower → round-trip wrapper
let buf2: RustBuffer =
<ToolRequestToolCall as FfiConverter<UniFfiTag>>::lower(ok_wrapper.clone());
let lifted_ok = <ToolRequestToolCall as FfiConverter<UniFfiTag>>::try_lift(buf2)
.expect("lift wrapper OK");
println!(
"ToolReq - Lifted wrapper (status: success): {:?}",
lifted_ok
);
assert_eq!(lifted_ok, ok_wrapper);
// ---------- status: error ----------
let err_call = ToolError::NotFound("no such function".into());
let err_wrapper = ToolRequestToolCall::from(Err(err_call.clone()));
let buf1: RustBuffer =
<ToolRequestToolCall as FfiConverter<UniFfiTag>>::lower(err_wrapper.clone());
let json_err: String =
<String as FfiConverter<UniFfiTag>>::try_lift(buf1).expect("lift String ERR");
println!("ToolReq - Lowered JSON (status: error): {:?}", json_err);
assert!(json_err.contains(r#""status":"error""#));
let buf2: RustBuffer =
<ToolRequestToolCall as FfiConverter<UniFfiTag>>::lower(err_wrapper.clone());
let lifted_err = <ToolRequestToolCall as FfiConverter<UniFfiTag>>::try_lift(buf2)
.expect("lift wrapper ERR");
println!("ToolReq - Lifted wrapper (status: error): {:?}", lifted_err);
match &*lifted_err {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("no such function"))
}
other => panic!("expected ExecutionError, got {:?}", other),
}
}
#[test]
fn ffi_roundtrip_tool_response_ok_and_err() {
// ---------- status: success ----------
let body = vec![Content::Text(TextContent {
text: "done".into(),
})];
let ok_wrapper = ToolResponseToolResult::from(Ok(body.clone()));
let buf1: RustBuffer =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::lower(ok_wrapper.clone());
let json_ok: String = <String as FfiConverter<UniFfiTag>>::try_lift(buf1).unwrap();
println!("ToolResp - Lowered JSON (status: success): {:?}", json_ok);
assert!(json_ok.contains(r#""status":"success""#));
let buf2: RustBuffer =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::lower(ok_wrapper.clone());
let lifted_ok =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::try_lift(buf2).unwrap();
println!(
"ToolResp - Lifted wrapper (status: success): {:?}",
lifted_ok
);
assert_eq!(lifted_ok, ok_wrapper);
// ---------- status: error ----------
let err_call = ToolError::InvalidParameters("bad params".into());
let err_wrapper = ToolResponseToolResult::from(Err(err_call.clone()));
let buf1: RustBuffer =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::lower(err_wrapper.clone());
let json_err: String = <String as FfiConverter<UniFfiTag>>::try_lift(buf1).unwrap();
println!("ToolResp - Lowered JSON (status: error): {:?}", json_err);
assert!(json_err.contains(r#""status":"error""#));
let buf2: RustBuffer =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::lower(err_wrapper.clone());
let lifted_err =
<ToolResponseToolResult as FfiConverter<UniFfiTag>>::try_lift(buf2).unwrap();
println!(
"ToolResp - Lifted wrapper (status: error): {:?}",
lifted_err
);
match &*lifted_err {
Err(ToolError::ExecutionError(msg)) => {
assert!(msg.contains("bad params"))
}
other => panic!("expected ExecutionError, got {:?}", other),
}
}
}
-284
View File
@@ -1,284 +0,0 @@
//! Messages which represent the content sent back and forth to LLM provider
//!
//! We use these messages in the agent code, and interfaces which interact with
//! the agent. That let's us reuse message histories across different interfaces.
//!
//! The content of the messages uses MCP types to avoid additional conversions
//! when interacting with MCP servers.
mod contents;
mod message_content;
mod tool_result_serde;
pub use contents::Contents;
pub use message_content::{
MessageContent, RedactedThinkingContent, ThinkingContent, ToolRequest, ToolRequestToolCall,
ToolResponse, ToolResponseToolResult,
};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use crate::types::core::Role;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
/// A message to or from an LLM
#[serde(rename_all = "camelCase")]
pub struct Message {
pub role: Role,
pub created: i64,
pub content: Contents,
}
impl Message {
pub fn new(role: Role) -> Self {
Self {
role,
created: Utc::now().timestamp_millis(),
content: Contents::default(),
}
}
/// Create a new user message with the current timestamp
pub fn user() -> Self {
Self::new(Role::User)
}
/// Create a new assistant message with the current timestamp
pub fn assistant() -> Self {
Self::new(Role::Assistant)
}
/// Add any item that implements Into<MessageContent> to the message
pub fn with_content(mut self, item: impl Into<MessageContent>) -> Self {
self.content.push(item);
self
}
/// Add text content to the message
pub fn with_text<S: Into<String>>(self, text: S) -> Self {
self.with_content(MessageContent::text(text))
}
/// Add image content to the message
pub fn with_image<S: Into<String>, T: Into<String>>(self, data: S, mime_type: T) -> Self {
self.with_content(MessageContent::image(data, mime_type))
}
/// Add a tool request to the message
pub fn with_tool_request<S: Into<String>, T: Into<ToolRequestToolCall>>(
self,
id: S,
tool_call: T,
) -> Self {
self.with_content(MessageContent::tool_request(id, tool_call.into()))
}
/// Add a tool response to the message
pub fn with_tool_response<S: Into<String>>(
self,
id: S,
result: ToolResponseToolResult,
) -> Self {
self.with_content(MessageContent::tool_response(id, result))
}
/// Add thinking content to the message
pub fn with_thinking<S1: Into<String>, S2: Into<String>>(
self,
thinking: S1,
signature: S2,
) -> Self {
self.with_content(MessageContent::thinking(thinking, signature))
}
/// Add redacted thinking content to the message
pub fn with_redacted_thinking<S: Into<String>>(self, data: S) -> Self {
self.with_content(MessageContent::redacted_thinking(data))
}
/// Check if the message is a tool call
pub fn contains_tool_call(&self) -> bool {
self.content.any_is(MessageContent::is_tool_request)
}
/// Check if the message is a tool response
pub fn contains_tool_response(&self) -> bool {
self.content.any_is(MessageContent::is_tool_response)
}
/// Check if the message contains only text content
pub fn has_only_text_content(&self) -> bool {
self.content.all_are(MessageContent::is_text)
}
/// Retrieves all tool `id` from ToolRequest messages
pub fn tool_request_ids(&self) -> HashSet<&str> {
self.content
.iter()
.filter_map(MessageContent::as_tool_request_id)
.collect()
}
/// Retrieves all tool `id` from ToolResponse messages
pub fn tool_response_ids(&self) -> HashSet<&str> {
self.content
.iter()
.filter_map(MessageContent::as_tool_response_id)
.collect()
}
/// Retrieves all tool `id` from the message
pub fn tool_ids(&self) -> HashSet<&str> {
self.tool_request_ids()
.into_iter()
.chain(self.tool_response_ids())
.collect()
}
}
#[cfg(test)]
mod tests {
use serde_json::{json, Value};
use super::*;
use crate::types::core::{ToolCall, ToolError};
#[test]
fn test_message_serialization() {
let message = Message::assistant()
.with_text("Hello, I'll help you with that.")
.with_tool_request(
"tool123",
Ok(ToolCall::new("test_tool", json!({"param": "value"}))),
);
let json_str = serde_json::to_string_pretty(&message).unwrap();
println!("Serialized message: {}", json_str);
// Parse back to Value to check structure
let value: Value = serde_json::from_str(&json_str).unwrap();
println!(
"Read back serialized message: {}",
serde_json::to_string_pretty(&value).unwrap()
);
// Check top-level fields
assert_eq!(value["role"], "assistant");
assert!(value["created"].is_i64());
assert!(value["content"].is_array());
// Check content items
let content = &value["content"];
// First item should be text
assert_eq!(content[0]["type"], "text");
assert_eq!(content[0]["text"], "Hello, I'll help you with that.");
// Second item should be toolRequest
assert_eq!(content[1]["type"], "toolReq");
assert_eq!(content[1]["id"], "tool123");
// Check tool_call serialization
assert_eq!(content[1]["toolCall"]["status"], "success");
assert_eq!(content[1]["toolCall"]["value"]["name"], "test_tool");
assert_eq!(
content[1]["toolCall"]["value"]["arguments"]["param"],
"value"
);
}
#[test]
fn test_error_serialization() {
let message = Message::assistant().with_tool_request(
"tool123",
Err(ToolError::ExecutionError(
"Something went wrong".to_string(),
)),
);
let json_str = serde_json::to_string_pretty(&message).unwrap();
println!("Serialized error: {}", json_str);
// Parse back to Value to check structure
let value: Value = serde_json::from_str(&json_str).unwrap();
// Check tool_call serialization with error
let tool_call = &value["content"][0]["toolCall"];
assert_eq!(tool_call["status"], "error");
assert_eq!(tool_call["error"], "Execution failed: Something went wrong");
}
#[test]
fn test_deserialization() {
// Create a JSON string with our new format
let json_str = r#"{
"role": "assistant",
"created": 1740171566,
"content": [
{
"type": "text",
"text": "I'll help you with that."
},
{
"type": "toolReq",
"id": "tool123",
"toolCall": {
"status": "success",
"value": {
"name": "test_tool",
"arguments": {"param": "value"},
"needsApproval": false
}
}
}
]
}"#;
let message: Message = serde_json::from_str(json_str).unwrap();
assert_eq!(message.role, Role::Assistant);
assert_eq!(message.created, 1740171566);
assert_eq!(message.content.len(), 2);
// Check first content item
if let MessageContent::Text(text) = &message.content[0] {
assert_eq!(text.text, "I'll help you with that.");
} else {
panic!("Expected Text content");
}
// Check second content item
if let MessageContent::ToolReq(req) = &message.content[1] {
assert_eq!(req.id, "tool123");
if let Ok(tool_call) = req.tool_call.as_result() {
assert_eq!(tool_call.name, "test_tool");
assert_eq!(tool_call.arguments, json!({"param": "value"}));
} else {
panic!("Expected successful tool call");
}
} else {
panic!("Expected ToolRequest content");
}
}
#[test]
fn test_message_with_text() {
let message = Message::user().with_text("Hello");
assert_eq!(message.content.concat_text_str(), "Hello");
}
#[test]
fn test_message_with_tool_request() {
let tool_call = Ok(ToolCall::new("test_tool", json!({})));
let message = Message::assistant().with_tool_request("req1", tool_call);
assert!(message.contains_tool_call());
assert!(!message.contains_tool_response());
let ids = message.tool_ids();
assert_eq!(ids.len(), 1);
assert!(ids.contains("req1"));
}
}
@@ -1,64 +0,0 @@
use serde::{ser::SerializeStruct, Deserialize, Deserializer, Serialize, Serializer};
use crate::types::core::{ToolError, ToolResult};
pub fn serialize<T, S>(value: &ToolResult<T>, serializer: S) -> Result<S::Ok, S::Error>
where
T: Serialize,
S: Serializer,
{
match value {
Ok(val) => {
let mut state = serializer.serialize_struct("ToolResult", 2)?;
state.serialize_field("status", "success")?;
state.serialize_field("value", val)?;
state.end()
}
Err(err) => {
let mut state = serializer.serialize_struct("ToolResult", 2)?;
state.serialize_field("status", "error")?;
state.serialize_field("error", &err.to_string())?;
state.end()
}
}
}
// For deserialization, let's use a simpler approach that works with the format we're serializing to
pub fn deserialize<'de, T, D>(deserializer: D) -> Result<ToolResult<T>, D::Error>
where
T: Deserialize<'de>,
D: Deserializer<'de>,
{
// Define a helper enum to handle the two possible formats
#[derive(Deserialize)]
#[serde(untagged)]
enum ResultFormat<T> {
Success { status: String, value: T },
Error { status: String, error: String },
}
let format = ResultFormat::deserialize(deserializer)?;
match format {
ResultFormat::Success { status, value } => {
if status == "success" {
Ok(Ok(value))
} else {
Err(serde::de::Error::custom(format!(
"Expected status 'success', got '{}'",
status
)))
}
}
ResultFormat::Error { status, error } => {
if status == "error" {
Ok(Err(ToolError::ExecutionError(error)))
} else {
Err(serde::de::Error::custom(format!(
"Expected status 'error', got '{}'",
status
)))
}
}
}
}
-119
View File
@@ -1,119 +0,0 @@
use serde::{Deserialize, Serialize};
const DEFAULT_CONTEXT_LIMIT: u32 = 128_000;
/// Configuration for model-specific settings and limits
#[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
pub struct ModelConfig {
/// The name of the model to use
pub model_name: String,
/// Optional explicit context limit that overrides any defaults
pub context_limit: Option<u32>,
/// Optional temperature setting (0.0 - 1.0)
pub temperature: Option<f32>,
/// Optional maximum tokens to generate
pub max_tokens: Option<i32>,
}
impl ModelConfig {
/// Create a new ModelConfig with the specified model name
///
/// The context limit is set with the following precedence:
/// 1. Explicit context_limit if provided in config
/// 2. Model-specific default based on model name
/// 3. Global default (128_000) (in get_context_limit)
pub fn new(model_name: String) -> Self {
let context_limit = Self::get_model_specific_limit(&model_name);
Self {
model_name,
context_limit,
temperature: None,
max_tokens: None,
}
}
/// Get model-specific context limit based on model name
fn get_model_specific_limit(model_name: &str) -> Option<u32> {
// Implement some sensible defaults
match model_name {
// OpenAI models, https://platform.openai.com/docs/models#models-overview
name if name.contains("gpt-4o") => Some(128_000),
name if name.contains("gpt-4-turbo") => Some(128_000),
// Anthropic models, https://docs.anthropic.com/en/docs/about-claude/models
name if name.contains("claude-3") => Some(200_000),
name if name.contains("claude-4") => Some(200_000),
// Meta Llama models, https://github.com/meta-llama/llama-models/tree/main?tab=readme-ov-file#llama-models-1
name if name.contains("llama3.2") => Some(128_000),
name if name.contains("llama3.3") => Some(128_000),
_ => None,
}
}
/// Set an explicit context limit
pub fn with_context_limit(mut self, limit: Option<u32>) -> Self {
// Default is None and therefore DEFAULT_CONTEXT_LIMIT, only set
// if input is Some to allow passing through with_context_limit in
// configuration cases
if limit.is_some() {
self.context_limit = limit;
}
self
}
/// Set the temperature
pub fn with_temperature(mut self, temp: Option<f32>) -> Self {
self.temperature = temp;
self
}
/// Set the max tokens
pub fn with_max_tokens(mut self, tokens: Option<i32>) -> Self {
self.max_tokens = tokens;
self
}
/// Get the context_limit for the current model
/// If none are defined, use the DEFAULT_CONTEXT_LIMIT
pub fn context_limit(&self) -> u32 {
self.context_limit.unwrap_or(DEFAULT_CONTEXT_LIMIT)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_model_config_context_limits() {
// Test explicit limit
let config =
ModelConfig::new("claude-3-opus".to_string()).with_context_limit(Some(150_000));
assert_eq!(config.context_limit(), 150_000);
// Test model-specific defaults
let config = ModelConfig::new("claude-3-opus".to_string());
assert_eq!(config.context_limit(), 200_000);
let config = ModelConfig::new("gpt-4-turbo".to_string());
assert_eq!(config.context_limit(), 128_000);
// Test fallback to default
let config = ModelConfig::new("unknown-model".to_string());
assert_eq!(config.context_limit(), DEFAULT_CONTEXT_LIMIT);
}
#[test]
fn test_model_config_settings() {
let config = ModelConfig::new("test-model".to_string())
.with_temperature(Some(0.7))
.with_max_tokens(Some(1000))
.with_context_limit(Some(50_000));
assert_eq!(config.temperature, Some(0.7));
assert_eq!(config.max_tokens, Some(1000));
assert_eq!(config.context_limit, Some(50_000));
}
}
-115
View File
@@ -1,115 +0,0 @@
use std::{
path::PathBuf,
sync::{Arc, RwLock},
};
use include_dir::{include_dir, Dir};
use minijinja::{Environment, Error as MiniJinjaError, Value as MJValue};
use once_cell::sync::Lazy;
use serde::Serialize;
/// This directory will be embedded into the final binary.
/// Typically used to store "core" or "system" prompts.
static CORE_PROMPTS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts");
/// A global MiniJinja environment storing the "core" prompts.
///
/// - Loaded at startup from the `CORE_PROMPTS_DIR`.
/// - Ideal for "system" templates that don't change often.
/// - *Not* used for extension prompts (which are ephemeral).
static GLOBAL_ENV: Lazy<Arc<RwLock<Environment<'static>>>> = Lazy::new(|| {
let mut env = Environment::new();
// Pre-load all core templates from the embedded dir.
for file in CORE_PROMPTS_DIR.files() {
let name = file.path().to_string_lossy().to_string();
let source = String::from_utf8_lossy(file.contents()).to_string();
// Since we're using 'static lifetime for the Environment, we need to ensure
// the strings we add as templates live for the entire program duration.
// We can achieve this by leaking the strings (acceptable for initialization).
let static_name: &'static str = Box::leak(name.into_boxed_str());
let static_source: &'static str = Box::leak(source.into_boxed_str());
if let Err(e) = env.add_template(static_name, static_source) {
println!("Failed to add template {}: {}", static_name, e);
}
}
Arc::new(RwLock::new(env))
});
/// Renders a prompt from the global environment by name.
///
/// # Arguments
/// * `template_name` - The name of the template (usually the file path or a custom ID).
/// * `context_data` - Data to be inserted into the template (must be `Serialize`).
pub fn render_global_template<T: Serialize>(
template_name: &str,
context_data: &T,
) -> Result<String, MiniJinjaError> {
let env = GLOBAL_ENV.read().expect("GLOBAL_ENV lock poisoned");
let tmpl = env.get_template(template_name)?;
let ctx = MJValue::from_serialize(context_data);
let rendered = tmpl.render(ctx)?;
Ok(rendered.trim().to_string())
}
/// Renders a file from `CORE_PROMPTS_DIR` within the global environment.
///
/// # Arguments
/// * `template_file` - The file path within the embedded directory (e.g. "system.md").
/// * `context_data` - Data to be inserted into the template (must be `Serialize`).
///
/// This function **assumes** the file is already in `CORE_PROMPTS_DIR`. If it wasn't
/// added to the global environment at startup (due to parse errors, etc.), this will error out.
pub fn render_global_file<T: Serialize>(
template_file: impl Into<PathBuf>,
context_data: &T,
) -> Result<String, MiniJinjaError> {
let file_path = template_file.into();
let template_name = file_path.to_string_lossy().to_string();
render_global_template(&template_name, context_data)
}
#[cfg(test)]
mod tests {
use super::*;
/// For convenience in tests, define a small struct or use a HashMap to provide context.
#[derive(Serialize)]
struct TestContext {
name: String,
age: u32,
}
#[test]
fn test_global_file_render() {
// "mock.md" should exist in the embedded CORE_PROMPTS_DIR
// and have placeholders for `name` and `age`.
let context = TestContext {
name: "Alice".to_string(),
age: 30,
};
let result = render_global_file("mock.md", &context).unwrap();
// Assume mock.md content is something like:
// "This prompt is only used for testing.\n\nHello, {{ name }}! You are {{ age }} years old."
assert_eq!(
result,
"This prompt is only used for testing.\n\nHello, Alice! You are 30 years old."
);
}
#[test]
fn test_global_file_not_found() {
let context = TestContext {
name: "Unused".to_string(),
age: 99,
};
let result = render_global_file("non_existent.md", &context);
assert!(result.is_err(), "Should fail because file is missing");
}
}
-3
View File
@@ -1,3 +0,0 @@
This prompt is only used for testing.
Hello, {{ name }}! You are {{ age }} years old.
-34
View File
@@ -1,34 +0,0 @@
{{system_preamble}}
The current date is {{current_date}}.
Goose uses LLM providers with tool calling capability. You can be used with different language models (gpt-4o, claude-3.5-sonnet, o1, llama-3.2, deepseek-r1, etc).
These models have varying knowledge cut-off dates depending on when they were trained, but typically it's between 5-10 months prior to the current date.
# Extensions
Extensions allow other applications to provide context to Goose. Extensions connect Goose to different data sources and tools.
{% if (extensions is defined) and extensions %}
Because you dynamically load extensions, your conversation history may refer
to interactions with extensions that are not currently active. The currently
active extensions are below. Each of these extensions provides tools that are
in your tool specification.
{% for extension in extensions %}
## {{extension.name}}
{% if extension.instructions %}### Instructions
{{extension.instructions}}{% endif %}
{% endfor %}
{% else %}
No extensions are defined. You should let the user know that they should add extensions.{% endif %}
# Response Guidelines
- Use Markdown formatting for all responses.
- Follow best practices for Markdown, including:
- Using headers for organization.
- Bullet points for lists.
- Links formatted correctly, either as linked text (e.g., [this is linked text](https://example.com)) or automatic links using angle brackets (e.g., <http://example.com/>).
- For code examples, use fenced code blocks by placing triple backticks (` ``` `) before and after the code. Include the language identifier after the opening backticks (e.g., ` ```python `) to enable syntax highlighting.
- Ensure clarity, conciseness, and proper formatting to enhance readability and usability.
-135
View File
@@ -1,135 +0,0 @@
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use super::errors::ProviderError;
use crate::{message::Message, types::core::Tool};
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize, uniffi::Record)]
pub struct Usage {
pub input_tokens: Option<i32>,
pub output_tokens: Option<i32>,
pub total_tokens: Option<i32>,
}
impl Usage {
pub fn new(
input_tokens: Option<i32>,
output_tokens: Option<i32>,
total_tokens: Option<i32>,
) -> Self {
Self {
input_tokens,
output_tokens,
total_tokens,
}
}
}
#[derive(Debug, Clone, uniffi::Record)]
pub struct ProviderCompleteResponse {
pub message: Message,
pub model: String,
pub usage: Usage,
}
impl ProviderCompleteResponse {
pub fn new(message: Message, model: String, usage: Usage) -> Self {
Self {
message,
model,
usage,
}
}
}
/// Response from a structuredextraction call
#[derive(Debug, Clone, uniffi::Record)]
pub struct ProviderExtractResponse {
/// The extracted JSON object
pub data: serde_json::Value,
/// Which model produced it
pub model: String,
/// Token usage stats
pub usage: Usage,
}
impl ProviderExtractResponse {
pub fn new(data: serde_json::Value, model: String, usage: Usage) -> Self {
Self { data, model, usage }
}
}
/// Base trait for AI providers (OpenAI, Anthropic, etc)
#[async_trait]
pub trait Provider: Send + Sync {
/// Generate the next message using the configured model and other parameters
///
/// # Arguments
/// * `system` - The system prompt that guides the model's behavior
/// * `messages` - The conversation history as a sequence of messages
/// * `tools` - Optional list of tools the model can use
/// * `request_id` - Optional request ID (only used by some providers like Databricks)
///
/// # Returns
/// A tuple containing the model's response message and provider usage statistics
///
/// # Errors
/// ProviderError
/// - It's important to raise ContextLengthExceeded correctly since agent handles it
async fn complete(
&self,
system: &str,
messages: &[Message],
tools: &[Tool],
request_id: Option<&str>,
) -> Result<ProviderCompleteResponse, ProviderError>;
/// Structured extraction: always JSONSchema
///
/// # Arguments
/// * `system` system prompt guiding the extraction task
/// * `messages` conversation history
/// * `schema` a JSONSchema for the expected output.
/// Will set strict=true for OpenAI & Databricks.
/// * `request_id` - Optional request ID (only used by some providers like Databricks)
///
/// # Returns
/// A `ProviderExtractResponse` whose `data` is a JSON object matching `schema`.
///
/// # Errors
/// * `ProviderError::ContextLengthExceeded` if the prompt is too large
/// * other `ProviderError` variants for API/network failures
async fn extract(
&self,
system: &str,
messages: &[Message],
schema: &serde_json::Value,
request_id: Option<&str>,
) -> Result<ProviderExtractResponse, ProviderError>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_usage_creation() {
let usage = Usage::new(Some(10), Some(20), Some(30));
assert_eq!(usage.input_tokens, Some(10));
assert_eq!(usage.output_tokens, Some(20));
assert_eq!(usage.total_tokens, Some(30));
}
#[test]
fn test_provider_complete_response_creation() {
let message = Message::user().with_text("Hello, world!");
let usage = Usage::new(Some(10), Some(20), Some(30));
let response =
ProviderCompleteResponse::new(message.clone(), "test_model".to_string(), usage.clone());
assert_eq!(response.message, message);
assert_eq!(response.model, "test_model");
assert_eq!(response.usage, usage);
}
}
@@ -1,327 +0,0 @@
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use url::Url;
use super::{
errors::ProviderError,
formats::databricks::{create_request, get_usage, response_to_message},
utils::{get_env, get_model, ImageFormat},
};
use crate::{
message::Message,
model::ModelConfig,
providers::{Provider, ProviderCompleteResponse, ProviderExtractResponse, Usage},
types::core::Tool,
};
pub const DATABRICKS_DEFAULT_MODEL: &str = "databricks-claude-3-7-sonnet";
// Databricks can passthrough to a wide range of models, we only provide the default
pub const _DATABRICKS_KNOWN_MODELS: &[&str] = &[
"databricks-meta-llama-3-3-70b-instruct",
"databricks-claude-3-7-sonnet",
];
fn default_timeout() -> u64 {
60
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabricksProviderConfig {
pub host: String,
pub token: String,
#[serde(default)]
pub image_format: ImageFormat,
#[serde(default = "default_timeout")]
pub timeout: u64, // timeout in seconds
}
impl DatabricksProviderConfig {
pub fn new(host: String, token: String) -> Self {
Self {
host,
token,
image_format: ImageFormat::OpenAi,
timeout: default_timeout(),
}
}
pub fn from_env() -> Self {
let host = get_env("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST");
let token = get_env("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN");
Self::new(host, token)
}
}
#[derive(Debug)]
pub struct DatabricksProvider {
config: DatabricksProviderConfig,
model: ModelConfig,
client: Client,
}
impl DatabricksProvider {
pub fn from_env(model: ModelConfig) -> Self {
let config = DatabricksProviderConfig::from_env();
DatabricksProvider::from_config(config, model)
.expect("Failed to initialize DatabricksProvider")
}
}
impl Default for DatabricksProvider {
fn default() -> Self {
let config = DatabricksProviderConfig::from_env();
let model = ModelConfig::new(DATABRICKS_DEFAULT_MODEL.to_string());
DatabricksProvider::from_config(config, model)
.expect("Failed to initialize DatabricksProvider")
}
}
impl DatabricksProvider {
pub fn from_config(config: DatabricksProviderConfig, model: ModelConfig) -> Result<Self> {
let client = Client::builder()
.timeout(Duration::from_secs(config.timeout))
.build()?;
Ok(Self {
config,
model,
client,
})
}
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = Url::parse(&self.config.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let path = format!("serving-endpoints/{}/invocations", self.model.model_name);
let url = base_url.join(&path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let auth_header = format!("Bearer {}", &self.config.token);
let response = self
.client
.post(url)
.header("Authorization", auth_header)
.json(&payload)
.send()
.await?;
let status = response.status();
let payload: Option<Value> = response.json().await.ok();
match status {
StatusCode::OK => payload.ok_or_else(|| {
ProviderError::RequestFailed("Response body is not valid JSON".to_string())
}),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
Err(ProviderError::Authentication(format!(
"Authentication failed. Please ensure your API keys are valid and have the required permissions. \
Status: {}. Response: {:?}",
status, payload
)))
}
StatusCode::BAD_REQUEST => {
// Databricks provides a generic 'error' but also includes 'external_model_message' which is provider specific
// We try to extract the error message from the payload and check for phrases that indicate context length exceeded
let payload_str = serde_json::to_string(&payload)
.unwrap_or_default()
.to_lowercase();
let check_phrases = [
"too long",
"context length",
"context_length_exceeded",
"reduce the length",
"token count",
"exceeds",
"exceed context limit",
"input length",
"max_tokens",
"decrease input length",
"context limit",
];
if check_phrases.iter().any(|c| payload_str.contains(c)) {
return Err(ProviderError::ContextLengthExceeded(payload_str));
}
let mut error_msg = "Unknown error".to_string();
if let Some(payload) = &payload {
// try to convert message to string, if that fails use external_model_message
error_msg = payload
.get("message")
.and_then(|m| m.as_str())
.or_else(|| {
payload
.get("external_model_message")
.and_then(|ext| ext.get("message"))
.and_then(|m| m.as_str())
})
.unwrap_or("Unknown error")
.to_string();
}
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status, payload
)
);
Err(ProviderError::RequestFailed(format!(
"Request failed with status: {}. Message: {}",
status, error_msg
)))
}
StatusCode::TOO_MANY_REQUESTS => {
Err(ProviderError::RateLimitExceeded(format!("{:?}", payload)))
}
StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => {
Err(ProviderError::ServerError(format!("{:?}", payload)))
}
_ => {
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status, payload
)
);
Err(ProviderError::RequestFailed(format!(
"Request failed with status: {}",
status
)))
}
}
}
}
#[async_trait]
impl Provider for DatabricksProvider {
#[tracing::instrument(
skip(self, system, messages, tools),
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
)]
async fn complete(
&self,
system: &str,
messages: &[Message],
tools: &[Tool],
request_id: Option<&str>,
) -> Result<ProviderCompleteResponse, ProviderError> {
let mut payload = create_request(
&self.model,
system,
messages,
tools,
&self.config.image_format,
)?;
// Remove the model key which is part of the url with databricks
payload
.as_object_mut()
.expect("payload should have model key")
.remove("model");
// Add client_request_id if provided
if let Some(req_id) = request_id {
payload
.as_object_mut()
.expect("payload should be an object")
.insert(
"client_request_id".to_string(),
serde_json::Value::String(req_id.to_string()),
);
}
let response = self.post(payload.clone()).await?;
// Parse response
let message = response_to_message(response.clone())?;
let usage = match get_usage(&response) {
Ok(usage) => usage,
Err(ProviderError::UsageError(e)) => {
tracing::debug!("Failed to get usage data: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let model = get_model(&response);
super::utils::emit_debug_trace(&self.model, &payload, &response, &usage);
Ok(ProviderCompleteResponse::new(message, model, usage))
}
async fn extract(
&self,
system: &str,
messages: &[Message],
schema: &Value,
request_id: Option<&str>,
) -> Result<ProviderExtractResponse, ProviderError> {
// 1. Build base payload (no tools)
let mut payload = create_request(&self.model, system, messages, &[], &ImageFormat::OpenAi)?;
// 2. Inject strict JSONSchema wrapper
payload
.as_object_mut()
.expect("payload must be an object")
.insert(
"response_format".to_string(),
json!({
"type": "json_schema",
"json_schema": {
"name": "extraction",
"schema": schema,
"strict": true
}
}),
);
// Add client_request_id if provided
if let Some(req_id) = request_id {
payload
.as_object_mut()
.expect("payload should be an object")
.insert(
"client_request_id".to_string(),
serde_json::Value::String(req_id.to_string()),
);
}
// 3. Call OpenAI
let response = self.post(payload.clone()).await?;
// 4. Extract the assistants `content` and parse it into JSON
let msg = &response["choices"][0]["message"];
let raw = msg.get("content").cloned().ok_or_else(|| {
ProviderError::ResponseParseError("Missing content in extract response".into())
})?;
let data = match raw {
Value::String(s) => serde_json::from_str(&s)
.map_err(|e| ProviderError::ResponseParseError(format!("Invalid JSON: {}", e)))?,
Value::Object(_) | Value::Array(_) => raw,
other => {
return Err(ProviderError::ResponseParseError(format!(
"Unexpected content type: {:?}",
other
)))
}
};
// 5. Gather usage & model info
let usage = match get_usage(&response) {
Ok(u) => u,
Err(ProviderError::UsageError(e)) => {
tracing::debug!("Failed to get usage in extract: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let model = get_model(&response);
Ok(ProviderExtractResponse::new(data, model, usage))
}
}
-144
View File
@@ -1,144 +0,0 @@
use thiserror::Error;
#[derive(Error, Debug, uniffi::Error)]
pub enum ProviderError {
#[error("Authentication error: {0}")]
Authentication(String),
#[error("Context length exceeded: {0}")]
ContextLengthExceeded(String),
#[error("Rate limit exceeded: {0}")]
RateLimitExceeded(String),
#[error("Server error: {0}")]
ServerError(String),
#[error("Request failed: {0}")]
RequestFailed(String),
#[error("Execution error: {0}")]
ExecutionError(String),
#[error("Usage data error: {0}")]
UsageError(String),
#[error("Invalid response: {0}")]
ResponseParseError(String),
}
impl From<anyhow::Error> for ProviderError {
fn from(error: anyhow::Error) -> Self {
ProviderError::ExecutionError(error.to_string())
}
}
impl From<reqwest::Error> for ProviderError {
fn from(error: reqwest::Error) -> Self {
ProviderError::ExecutionError(error.to_string())
}
}
#[derive(serde::Deserialize, Debug)]
pub struct OpenAIError {
#[serde(deserialize_with = "code_as_string")]
pub code: Option<String>,
pub message: Option<String>,
#[serde(rename = "type")]
pub error_type: Option<String>,
}
fn code_as_string<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
D: serde::Deserializer<'de>,
{
use std::fmt;
use serde::de::{self, Visitor};
struct CodeVisitor;
impl<'de> Visitor<'de> for CodeVisitor {
type Value = Option<String>;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a string, a number, null, or none for the code field")
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(Some(value.to_string()))
}
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(Some(value.to_string()))
}
fn visit_none<E>(self) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(None)
}
fn visit_unit<E>(self) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(None)
}
fn visit_some<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_any(CodeVisitor)
}
}
deserializer.deserialize_option(CodeVisitor)
}
impl OpenAIError {
pub fn is_context_length_exceeded(&self) -> bool {
if let Some(code) = &self.code {
code == "context_length_exceeded" || code == "string_above_max_length"
} else {
false
}
}
}
impl std::fmt::Display for OpenAIError {
/// Format the error for display.
/// E.g. {"message": "Invalid API key", "code": "invalid_api_key", "type": "client_error"}
/// would be formatted as "Invalid API key (code: invalid_api_key, type: client_error)"
/// and {"message": "Foo"} as just "Foo", etc.
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(message) = &self.message {
write!(f, "{}", message)?;
}
let mut in_parenthesis = false;
if let Some(code) = &self.code {
write!(f, " (code: {}", code)?;
in_parenthesis = true;
}
if let Some(typ) = &self.error_type {
if in_parenthesis {
write!(f, ", type: {}", typ)?;
} else {
write!(f, " (type: {}", typ)?;
in_parenthesis = true;
}
}
if in_parenthesis {
write!(f, ")")?;
}
Ok(())
}
}
-29
View File
@@ -1,29 +0,0 @@
use std::sync::Arc;
use anyhow::Result;
use super::{
base::Provider,
databricks::{DatabricksProvider, DatabricksProviderConfig},
openai::{OpenAiProvider, OpenAiProviderConfig},
};
use crate::model::ModelConfig;
pub fn create(
name: &str,
provider_config: serde_json::Value,
model: ModelConfig,
) -> Result<Arc<dyn Provider>> {
// We use Arc instead of Box to be able to clone for multiple async tasks
match name {
"openai" => {
let config: OpenAiProviderConfig = serde_json::from_value(provider_config)?;
Ok(Arc::new(OpenAiProvider::from_config(config, model)?))
}
"databricks" => {
let config: DatabricksProviderConfig = serde_json::from_value(provider_config)?;
Ok(Arc::new(DatabricksProvider::from_config(config, model)?))
}
_ => Err(anyhow::anyhow!("Unknown provider: {}", name)),
}
}
File diff suppressed because it is too large Load Diff
@@ -1,2 +0,0 @@
pub mod databricks;
pub mod openai;
@@ -1,846 +0,0 @@
use anyhow::{anyhow, Error};
use serde_json::{json, Value};
use crate::{
message::{Message, MessageContent},
model::ModelConfig,
providers::{
base::Usage,
errors::ProviderError,
utils::{convert_image, is_valid_function_name, sanitize_function_name, ImageFormat},
},
types::core::{Content, Role, Tool, ToolCall, ToolError},
};
/// Convert internal Message format to OpenAI's API message specification
/// some openai compatible endpoints use the anthropic image spec at the content level
/// even though the message structure is otherwise following openai, the enum switches this
pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Value> {
let mut messages_spec = Vec::new();
for message in messages {
let mut converted = json!({
"role": message.role
});
let mut output = Vec::new();
for content in message.content.iter() {
match content {
MessageContent::Text(text) => {
if !text.text.is_empty() {
converted["content"] = json!(text.text);
}
}
MessageContent::Image(image) => {
// Handle direct image content
converted["content"] = json!([convert_image(image, image_format)]);
}
MessageContent::Thinking(_) => {
// Thinking blocks are not directly used in OpenAI format
continue;
}
MessageContent::RedactedThinking(_) => {
// Redacted thinking blocks are not directly used in OpenAI format
continue;
}
MessageContent::ToolReq(request) => match &request.tool_call.as_result() {
Ok(tool_call) => {
let sanitized_name = sanitize_function_name(&tool_call.name);
let tool_calls = converted
.as_object_mut()
.unwrap()
.entry("tool_calls")
.or_insert(json!([]));
tool_calls.as_array_mut().unwrap().push(json!({
"id": request.id,
"type": "function",
"function": {
"name": sanitized_name,
"arguments": tool_call.arguments.to_string(),
}
}));
}
Err(e) => {
output.push(json!({
"role": "tool",
"content": format!("Error: {}", e),
"tool_call_id": request.id
}));
}
},
MessageContent::ToolResp(response) => {
match &response.tool_result.0 {
Ok(contents) => {
// Process all content, replacing images with placeholder text
let mut tool_content = Vec::new();
let mut image_messages = Vec::new();
for content in contents {
match content {
Content::Image(image) => {
// Add placeholder text in the tool response
tool_content.push(Content::text("This tool result included an image that is uploaded in the next message."));
// Create a separate image message
image_messages.push(json!({
"role": "user",
"content": [convert_image(image, image_format)]
}));
}
_ => {
tool_content.push(content.clone());
}
}
}
let tool_response_content: Value = json!(tool_content
.iter()
.map(|content| match content {
Content::Text(text) => text.text.clone(),
_ => String::new(),
})
.collect::<Vec<String>>()
.join(" "));
// First add the tool response with all content
output.push(json!({
"role": "tool",
"content": tool_response_content,
"tool_call_id": response.id
}));
// Then add any image messages that need to follow
output.extend(image_messages);
}
Err(e) => {
// A tool result error is shown as output so the model can interpret the error message
output.push(json!({
"role": "tool",
"content": format!("The tool call returned the following error:\n{}", e),
"tool_call_id": response.id
}));
}
}
}
}
}
if converted.get("content").is_some() || converted.get("tool_calls").is_some() {
output.insert(0, converted);
}
messages_spec.extend(output);
}
messages_spec
}
/// Convert internal Tool format to OpenAI's API tool specification
pub fn format_tools(tools: &[Tool]) -> anyhow::Result<Vec<Value>> {
let mut tool_names = std::collections::HashSet::new();
let mut result = Vec::new();
for tool in tools {
if !tool_names.insert(&tool.name) {
return Err(anyhow!("Duplicate tool name: {}", tool.name));
}
let mut description = tool.description.clone();
description.truncate(1024);
// OpenAI's tool description max str len is 1024
result.push(json!({
"type": "function",
"function": {
"name": tool.name,
"description": description,
"parameters": tool.input_schema,
}
}));
}
Ok(result)
}
/// Convert OpenAI's API response to internal Message format
pub fn response_to_message(response: Value) -> anyhow::Result<Message> {
let original = response["choices"][0]["message"].clone();
let mut content = Vec::new();
if let Some(text) = original.get("content") {
if let Some(text_str) = text.as_str() {
content.push(MessageContent::text(text_str));
}
}
if let Some(tool_calls) = original.get("tool_calls") {
if let Some(tool_calls_array) = tool_calls.as_array() {
for tool_call in tool_calls_array {
let id = tool_call["id"].as_str().unwrap_or_default().to_string();
let function_name = tool_call["function"]["name"]
.as_str()
.unwrap_or_default()
.to_string();
let mut arguments = tool_call["function"]["arguments"]
.as_str()
.unwrap_or_default()
.to_string();
// If arguments is empty, we will have invalid json parsing error later.
if arguments.is_empty() {
arguments = "{}".to_string();
}
if !is_valid_function_name(&function_name) {
let error = ToolError::NotFound(format!(
"The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+",
function_name
));
content.push(MessageContent::tool_request(id, Err(error).into()));
} else {
match serde_json::from_str::<Value>(&arguments) {
Ok(params) => {
content.push(MessageContent::tool_request(
id,
Ok(ToolCall::new(&function_name, params)).into(),
));
}
Err(e) => {
let error = ToolError::InvalidParameters(format!(
"Could not interpret tool use parameters for id {}: {}",
id, e
));
content.push(MessageContent::tool_request(id, Err(error).into()));
}
}
}
}
}
}
Ok(Message {
role: Role::Assistant,
created: chrono::Utc::now().timestamp_millis(),
content: content.into(),
})
}
pub fn get_usage(data: &Value) -> Result<Usage, ProviderError> {
let usage = data
.get("usage")
.ok_or_else(|| ProviderError::UsageError("No usage data in response".to_string()))?;
let input_tokens = usage
.get("prompt_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32);
let output_tokens = usage
.get("completion_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32);
let total_tokens = usage
.get("total_tokens")
.and_then(|v| v.as_i64())
.map(|v| v as i32)
.or_else(|| match (input_tokens, output_tokens) {
(Some(input), Some(output)) => Some(input + output),
_ => None,
});
Ok(Usage::new(input_tokens, output_tokens, total_tokens))
}
/// Validates and fixes tool schemas to ensure they have proper parameter structure.
/// If parameters exist, ensures they have properties and required fields, or removes parameters entirely.
pub fn validate_tool_schemas(tools: &mut [Value]) {
for tool in tools.iter_mut() {
if let Some(function) = tool.get_mut("function") {
if let Some(parameters) = function.get_mut("parameters") {
if parameters.is_object() {
ensure_valid_json_schema(parameters);
}
}
}
}
}
/// Ensures that the given JSON value follows the expected JSON Schema structure.
fn ensure_valid_json_schema(schema: &mut Value) {
if let Some(params_obj) = schema.as_object_mut() {
// Check if this is meant to be an object type schema
let is_object_type = params_obj
.get("type")
.and_then(|t| t.as_str())
.is_none_or(|t| t == "object"); // Default to true if no type is specified
// Only apply full schema validation to object types
if is_object_type {
// Ensure required fields exist with default values
params_obj.entry("properties").or_insert_with(|| json!({}));
params_obj.entry("required").or_insert_with(|| json!([]));
params_obj.entry("type").or_insert_with(|| json!("object"));
// Recursively validate properties if it exists
if let Some(properties) = params_obj.get_mut("properties") {
if let Some(properties_obj) = properties.as_object_mut() {
for (_key, prop) in properties_obj.iter_mut() {
if prop.is_object()
&& prop.get("type").and_then(|t| t.as_str()) == Some("object")
{
ensure_valid_json_schema(prop);
}
}
}
}
}
}
}
pub fn create_request(
model_config: &ModelConfig,
system: &str,
messages: &[Message],
tools: &[Tool],
image_format: &ImageFormat,
) -> anyhow::Result<Value, Error> {
if model_config.model_name.starts_with("o1-mini") {
return Err(anyhow!(
"o1-mini model is not currently supported since Goose uses tool calling and o1-mini does not support it. Please use o1 or o3 models instead."
));
}
let is_ox_model = model_config.model_name.starts_with("o");
// Only extract reasoning effort for O1/O3 models
let (model_name, reasoning_effort) = if is_ox_model {
let parts: Vec<&str> = model_config.model_name.split('-').collect();
let last_part = parts.last().unwrap();
match *last_part {
"low" | "medium" | "high" => {
let base_name = parts[..parts.len() - 1].join("-");
(base_name, Some(last_part.to_string()))
}
_ => (
model_config.model_name.to_string(),
Some("medium".to_string()),
),
}
} else {
// For non-O family models, use the model name as is and no reasoning effort
(model_config.model_name.to_string(), None)
};
let system_message = json!({
"role": if is_ox_model { "developer" } else { "system" },
"content": system
});
let messages_spec = format_messages(messages, image_format);
let mut tools_spec = if !tools.is_empty() {
format_tools(tools)?
} else {
vec![]
};
// Validate tool schemas
validate_tool_schemas(&mut tools_spec);
let mut messages_array = vec![system_message];
messages_array.extend(messages_spec);
let mut payload = json!({
"model": model_name,
"messages": messages_array
});
if let Some(effort) = reasoning_effort {
payload
.as_object_mut()
.unwrap()
.insert("reasoning_effort".to_string(), json!(effort));
}
if !tools_spec.is_empty() {
payload
.as_object_mut()
.unwrap()
.insert("tools".to_string(), json!(tools_spec));
}
// o1, o3 models currently don't support temperature
if !is_ox_model {
if let Some(temp) = model_config.temperature {
payload
.as_object_mut()
.unwrap()
.insert("temperature".to_string(), json!(temp));
}
}
// o1 models use max_completion_tokens instead of max_tokens
if let Some(tokens) = model_config.max_tokens {
let key = if is_ox_model {
"max_completion_tokens"
} else {
"max_tokens"
};
payload
.as_object_mut()
.unwrap()
.insert(key.to_string(), json!(tokens));
}
Ok(payload)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
use crate::types::core::Content;
#[test]
fn test_validate_tool_schemas() {
// Test case 1: Empty parameters object
// Input JSON with an incomplete parameters object
let mut actual = vec![json!({
"type": "function",
"function": {
"name": "test_func",
"description": "test description",
"parameters": {
"type": "object"
}
}
})];
// Run the function to validate and update schemas
validate_tool_schemas(&mut actual);
// Expected JSON after validation
let expected = vec![json!({
"type": "function",
"function": {
"name": "test_func",
"description": "test description",
"parameters": {
"type": "object",
"properties": {},
"required": []
}
}
})];
// Compare entire JSON structures instead of individual fields
assert_eq!(actual, expected);
// Test case 2: Missing type field
let mut tools = vec![json!({
"type": "function",
"function": {
"name": "test_func",
"description": "test description",
"parameters": {
"properties": {}
}
}
})];
validate_tool_schemas(&mut tools);
let params = tools[0]["function"]["parameters"].as_object().unwrap();
assert_eq!(params["type"], "object");
// Test case 3: Complete valid schema should remain unchanged
let original_schema = json!({
"type": "function",
"function": {
"name": "test_func",
"description": "test description",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "City and country"
}
},
"required": ["location"]
}
}
});
let mut tools = vec![original_schema.clone()];
validate_tool_schemas(&mut tools);
assert_eq!(tools[0], original_schema);
}
const OPENAI_TOOL_USE_RESPONSE: &str = r#"{
"choices": [{
"role": "assistant",
"message": {
"tool_calls": [{
"id": "1",
"function": {
"name": "example_fn",
"arguments": "{\"param\": \"value\"}"
}
}]
}
}],
"usage": {
"input_tokens": 10,
"output_tokens": 25,
"total_tokens": 35
}
}"#;
#[test]
fn test_format_messages() -> anyhow::Result<()> {
let message = Message::user().with_text("Hello");
let spec = format_messages(&[message], &ImageFormat::OpenAi);
assert_eq!(spec.len(), 1);
assert_eq!(spec[0]["role"], "user");
assert_eq!(spec[0]["content"], "Hello");
Ok(())
}
#[test]
fn test_format_tools() -> anyhow::Result<()> {
let tool = Tool::new(
"test_tool",
"A test tool",
json!({
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Test parameter"
}
},
"required": ["input"]
}),
);
let spec = format_tools(&[tool])?;
assert_eq!(spec.len(), 1);
assert_eq!(spec[0]["type"], "function");
assert_eq!(spec[0]["function"]["name"], "test_tool");
Ok(())
}
#[test]
fn test_format_messages_complex() -> anyhow::Result<()> {
let mut messages = vec![
Message::assistant().with_text("Hello!"),
Message::user().with_text("How are you?"),
Message::assistant().with_tool_request(
"tool1",
Ok(ToolCall::new("example", json!({"param1": "value1"}))),
),
];
// Get the ID from the tool request to use in the response
let tool_id = if let MessageContent::ToolReq(request) = &messages[2].content[0] {
request.id.clone()
} else {
panic!("should be tool request");
};
messages.push(
Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]).into()),
);
let spec = format_messages(&messages, &ImageFormat::OpenAi);
assert_eq!(spec.len(), 4);
assert_eq!(spec[0]["role"], "assistant");
assert_eq!(spec[0]["content"], "Hello!");
assert_eq!(spec[1]["role"], "user");
assert_eq!(spec[1]["content"], "How are you?");
assert_eq!(spec[2]["role"], "assistant");
assert!(spec[2]["tool_calls"].is_array());
assert_eq!(spec[3]["role"], "tool");
assert_eq!(spec[3]["content"], "Result");
assert_eq!(spec[3]["tool_call_id"], spec[2]["tool_calls"][0]["id"]);
Ok(())
}
#[test]
fn test_format_messages_multiple_content() -> anyhow::Result<()> {
let mut messages = vec![Message::assistant().with_tool_request(
"tool1",
Ok(ToolCall::new("example", json!({"param1": "value1"}))),
)];
// Get the ID from the tool request to use in the response
let tool_id = if let MessageContent::ToolReq(request) = &messages[0].content[0] {
request.id.clone()
} else {
panic!("should be tool request");
};
messages.push(
Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]).into()),
);
let spec = format_messages(&messages, &ImageFormat::OpenAi);
assert_eq!(spec.len(), 2);
assert_eq!(spec[0]["role"], "assistant");
assert!(spec[0]["tool_calls"].is_array());
assert_eq!(spec[1]["role"], "tool");
assert_eq!(spec[1]["content"], "Result");
assert_eq!(spec[1]["tool_call_id"], spec[0]["tool_calls"][0]["id"]);
Ok(())
}
#[test]
fn test_format_tools_duplicate() -> anyhow::Result<()> {
let tool1 = Tool::new(
"test_tool",
"Test tool",
json!({
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Test parameter"
}
},
"required": ["input"]
}),
);
let tool2 = Tool::new(
"test_tool",
"Test tool",
json!({
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Test parameter"
}
},
"required": ["input"]
}),
);
let result = format_tools(&[tool1, tool2]);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Duplicate tool name"));
Ok(())
}
#[test]
fn test_format_tools_empty() -> anyhow::Result<()> {
let spec = format_tools(&[])?;
assert!(spec.is_empty());
Ok(())
}
#[test]
fn test_response_to_message_text() -> anyhow::Result<()> {
let response = json!({
"choices": [{
"role": "assistant",
"message": {
"content": "Hello from John Cena!"
}
}],
"usage": {
"input_tokens": 10,
"output_tokens": 25,
"total_tokens": 35
}
});
let message = response_to_message(response)?;
assert_eq!(message.content.len(), 1);
if let MessageContent::Text(text) = &message.content[0] {
assert_eq!(text.text, "Hello from John Cena!");
} else {
panic!("Expected Text content");
}
assert!(matches!(message.role, Role::Assistant));
Ok(())
}
#[test]
fn test_response_to_message_valid_toolrequest() -> anyhow::Result<()> {
let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?;
let message = response_to_message(response)?;
assert_eq!(message.content.len(), 1);
if let MessageContent::ToolReq(request) = &message.content[0] {
let tool_call = request.tool_call.as_ref().unwrap();
assert_eq!(tool_call.name, "example_fn");
assert_eq!(tool_call.arguments, json!({"param": "value"}));
} else {
panic!("Expected ToolRequest content");
}
Ok(())
}
#[test]
fn test_response_to_message_invalid_func_name() -> anyhow::Result<()> {
let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?;
response["choices"][0]["message"]["tool_calls"][0]["function"]["name"] =
json!("invalid fn");
let message = response_to_message(response)?;
if let MessageContent::ToolReq(request) = &message.content[0] {
match &request.tool_call.as_result() {
Err(ToolError::NotFound(msg)) => {
assert!(msg.starts_with("The provided function name"));
}
_ => panic!("Expected ToolNotFound error"),
}
} else {
panic!("Expected ToolRequest content");
}
Ok(())
}
#[test]
fn test_response_to_message_json_decode_error() -> anyhow::Result<()> {
let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?;
response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] =
json!("invalid json {");
let message = response_to_message(response)?;
if let MessageContent::ToolReq(request) = &message.content[0] {
match &request.tool_call.as_result() {
Err(ToolError::InvalidParameters(msg)) => {
assert!(msg.starts_with("Could not interpret tool use parameters"));
}
_ => panic!("Expected InvalidParameters error"),
}
} else {
panic!("Expected ToolRequest content");
}
Ok(())
}
#[test]
fn test_response_to_message_empty_argument() -> anyhow::Result<()> {
let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?;
response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] =
serde_json::Value::String("".to_string());
let message = response_to_message(response)?;
if let MessageContent::ToolReq(request) = &message.content[0] {
let tool_call = request.tool_call.as_ref().unwrap();
assert_eq!(tool_call.name, "example_fn");
assert_eq!(tool_call.arguments, json!({}));
} else {
panic!("Expected ToolRequest content");
}
Ok(())
}
#[test]
fn test_create_request_gpt_4o() -> anyhow::Result<()> {
// Test default medium reasoning effort for O3 model
let model_config = ModelConfig {
model_name: "gpt-4o".to_string(),
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
let expected = json!({
"model": "gpt-4o",
"messages": [
{
"role": "system",
"content": "system"
}
],
"max_tokens": 1024
});
for (key, value) in expected.as_object().unwrap() {
assert_eq!(obj.get(key).unwrap(), value);
}
Ok(())
}
#[test]
fn test_create_request_o1_default() -> anyhow::Result<()> {
// Test default medium reasoning effort for O1 model
let model_config = ModelConfig {
model_name: "o1".to_string(),
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
let expected = json!({
"model": "o1",
"messages": [
{
"role": "developer",
"content": "system"
}
],
"reasoning_effort": "medium",
"max_completion_tokens": 1024
});
for (key, value) in expected.as_object().unwrap() {
assert_eq!(obj.get(key).unwrap(), value);
}
Ok(())
}
#[test]
fn test_create_request_o3_custom_reasoning_effort() -> anyhow::Result<()> {
// Test custom reasoning effort for O3 model
let model_config = ModelConfig {
model_name: "o3-mini-high".to_string(),
context_limit: Some(4096),
temperature: None,
max_tokens: Some(1024),
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
let obj = request.as_object().unwrap();
let expected = json!({
"model": "o3-mini",
"messages": [
{
"role": "developer",
"content": "system"
}
],
"reasoning_effort": "high",
"max_completion_tokens": 1024
});
for (key, value) in expected.as_object().unwrap() {
assert_eq!(obj.get(key).unwrap(), value);
}
Ok(())
}
}
-10
View File
@@ -1,10 +0,0 @@
pub mod base;
pub mod databricks;
pub mod errors;
mod factory;
pub mod formats;
pub mod openai;
pub mod utils;
pub use base::{Provider, ProviderCompleteResponse, ProviderExtractResponse, Usage};
pub use factory::create;
-233
View File
@@ -1,233 +0,0 @@
use std::{collections::HashMap, time::Duration};
use anyhow::Result;
use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use super::{
errors::ProviderError,
formats::openai::{create_request, get_usage, response_to_message},
utils::{emit_debug_trace, get_env, get_model, handle_response_openai_compat, ImageFormat},
};
use crate::{
message::Message,
model::ModelConfig,
providers::{Provider, ProviderCompleteResponse, ProviderExtractResponse, Usage},
types::core::Tool,
};
pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o";
pub const _OPEN_AI_KNOWN_MODELS: &[&str] = &["gpt-4o", "gpt-4.1", "o1", "o3", "o4-mini"];
fn default_timeout() -> u64 {
60
}
fn default_base_path() -> String {
"v1/chat/completions".to_string()
}
fn default_host() -> String {
"https://api.openai.com".to_string()
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OpenAiProviderConfig {
pub api_key: String,
#[serde(default = "default_host")]
pub host: String,
#[serde(default)]
pub organization: Option<String>,
#[serde(default = "default_base_path")]
pub base_path: String,
#[serde(default)]
pub project: Option<String>,
#[serde(default)]
pub custom_headers: Option<HashMap<String, String>>,
#[serde(default = "default_timeout")]
pub timeout: u64, // timeout in seconds
}
impl OpenAiProviderConfig {
pub fn new(api_key: String) -> Self {
Self {
api_key,
host: default_host(),
organization: None,
base_path: default_base_path(),
project: None,
custom_headers: None,
timeout: 600,
}
}
pub fn from_env() -> Self {
let api_key = get_env("OPENAI_API_KEY").expect("Missing OPENAI_API_KEY");
Self::new(api_key)
}
}
#[derive(Debug)]
pub struct OpenAiProvider {
config: OpenAiProviderConfig,
model: ModelConfig,
client: Client,
}
impl OpenAiProvider {
pub fn from_env(model: ModelConfig) -> Self {
let config = OpenAiProviderConfig::from_env();
OpenAiProvider::from_config(config, model).expect("Failed to initialize OpenAiProvider")
}
}
impl Default for OpenAiProvider {
fn default() -> Self {
let config = OpenAiProviderConfig::from_env();
let model = ModelConfig::new(OPEN_AI_DEFAULT_MODEL.to_string());
OpenAiProvider::from_config(config, model).expect("Failed to initialize OpenAiProvider")
}
}
impl OpenAiProvider {
pub fn from_config(config: OpenAiProviderConfig, model: ModelConfig) -> Result<Self> {
let client = Client::builder()
.timeout(Duration::from_secs(config.timeout))
.build()?;
Ok(Self {
config,
model,
client,
})
}
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let base_url = url::Url::parse(&self.config.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join(&self.config.base_path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let mut request = self
.client
.post(url)
.header("Authorization", format!("Bearer {}", self.config.api_key));
// Add organization header if present
if let Some(org) = &self.config.organization {
request = request.header("OpenAI-Organization", org);
}
// Add project header if present
if let Some(project) = &self.config.project {
request = request.header("OpenAI-Project", project);
}
if let Some(custom_headers) = &self.config.custom_headers {
for (key, value) in custom_headers {
request = request.header(key, value);
}
}
let response = request.json(&payload).send().await?;
handle_response_openai_compat(response).await
}
}
#[async_trait]
impl Provider for OpenAiProvider {
#[tracing::instrument(
skip(self, system, messages, tools),
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
)]
async fn complete(
&self,
system: &str,
messages: &[Message],
tools: &[Tool],
_request_id: Option<&str>, // OpenAI doesn't use request_id, so we ignore it
) -> Result<ProviderCompleteResponse, ProviderError> {
let payload = create_request(&self.model, system, messages, tools, &ImageFormat::OpenAi)?;
// Make request
let response = self.post(payload.clone()).await?;
// Parse response
let message = response_to_message(response.clone())?;
let usage = match get_usage(&response) {
Ok(usage) => usage,
Err(ProviderError::UsageError(e)) => {
tracing::debug!("Failed to get usage data: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let model = get_model(&response);
emit_debug_trace(&self.model, &payload, &response, &usage);
Ok(ProviderCompleteResponse::new(message, model, usage))
}
async fn extract(
&self,
system: &str,
messages: &[Message],
schema: &Value,
_request_id: Option<&str>, // OpenAI doesn't use request_id, so we ignore it
) -> Result<ProviderExtractResponse, ProviderError> {
// 1. Build base payload (no tools)
let mut payload = create_request(&self.model, system, messages, &[], &ImageFormat::OpenAi)?;
// 2. Inject strict JSONSchema wrapper
payload
.as_object_mut()
.expect("payload must be an object")
.insert(
"response_format".to_string(),
json!({
"type": "json_schema",
"json_schema": {
"name": "extraction",
"schema": schema,
"strict": true
}
}),
);
// 3. Call OpenAI
let response = self.post(payload.clone()).await?;
// 4. Extract the assistants `content` and parse it into JSON
let msg = &response["choices"][0]["message"];
let raw = msg.get("content").cloned().ok_or_else(|| {
ProviderError::ResponseParseError("Missing content in extract response".into())
})?;
let data = match raw {
Value::String(s) => serde_json::from_str(&s)
.map_err(|e| ProviderError::ResponseParseError(format!("Invalid JSON: {}", e)))?,
Value::Object(_) | Value::Array(_) => raw,
other => {
return Err(ProviderError::ResponseParseError(format!(
"Unexpected content type: {:?}",
other
)))
}
};
// 5. Gather usage & model info
let usage = match get_usage(&response) {
Ok(u) => u,
Err(ProviderError::UsageError(e)) => {
tracing::debug!("Failed to get usage in extract: {}", e);
Usage::default()
}
Err(e) => return Err(e),
};
let model = get_model(&response);
Ok(ProviderExtractResponse::new(data, model, usage))
}
}
-260
View File
@@ -1,260 +0,0 @@
use std::{env, io::Read, path::Path};
use anyhow::Result;
use base64::Engine;
use regex::Regex;
use reqwest::{Response, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::{from_value, json, Value};
use super::base::Usage;
use crate::{
model::ModelConfig,
providers::errors::{OpenAIError, ProviderError},
types::core::ImageContent,
};
#[derive(serde::Deserialize)]
struct OpenAIErrorResponse {
error: OpenAIError,
}
#[derive(Debug, Copy, Clone, Serialize, Deserialize, Default)]
pub enum ImageFormat {
#[default]
OpenAi,
Anthropic,
}
/// Timeout in seconds.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Timeout(u32);
impl Default for Timeout {
fn default() -> Self {
Timeout(60)
}
}
/// Convert an image content into an image json based on format
pub fn convert_image(image: &ImageContent, image_format: &ImageFormat) -> Value {
match image_format {
ImageFormat::OpenAi => json!({
"type": "image_url",
"image_url": {
"url": format!("data:{};base64,{}", image.mime_type, image.data)
}
}),
ImageFormat::Anthropic => json!({
"type": "image",
"source": {
"type": "base64",
"media_type": image.mime_type,
"data": image.data,
}
}),
}
}
/// Handle response from OpenAI compatible endpoints
/// Error codes: https://platform.openai.com/docs/guides/error-codes
/// Context window exceeded: https://community.openai.com/t/help-needed-tackling-context-length-limits-in-openai-models/617543
pub async fn handle_response_openai_compat(response: Response) -> Result<Value, ProviderError> {
let status = response.status();
// Try to parse the response body as JSON (if applicable)
let payload = match response.json::<Value>().await {
Ok(json) => json,
Err(e) => return Err(ProviderError::RequestFailed(e.to_string())),
};
match status {
StatusCode::OK => Ok(payload),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
Err(ProviderError::Authentication(format!(
"Authentication failed. Please ensure your API keys are valid and have the required permissions. \
Status: {}. Response: {:?}",
status, payload
)))
}
StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND => {
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status, payload
)
);
if let Ok(err_resp) = from_value::<OpenAIErrorResponse>(payload) {
let err = err_resp.error;
if err.is_context_length_exceeded() {
return Err(ProviderError::ContextLengthExceeded(
err.message.unwrap_or("Unknown error".to_string()),
));
}
return Err(ProviderError::RequestFailed(format!(
"{} (status {})",
err,
status.as_u16()
)));
}
Err(ProviderError::RequestFailed(format!(
"Unknown error (status {})",
status
)))
}
StatusCode::TOO_MANY_REQUESTS => {
Err(ProviderError::RateLimitExceeded(format!("{:?}", payload)))
}
StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => {
Err(ProviderError::ServerError(format!("{:?}", payload)))
}
_ => {
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status, payload
)
);
Err(ProviderError::RequestFailed(format!(
"Request failed with status: {}",
status
)))
}
}
}
/// Get a secret from environment variables. The secret is expected to be in JSON format.
pub fn get_env(key: &str) -> Result<String> {
// check environment variables (convert to uppercase)
let env_key = key.to_uppercase();
if let Ok(val) = env::var(&env_key) {
let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val));
Ok(serde_json::from_value(value)?)
} else {
Err(anyhow::anyhow!(
"Environment variable {} not found",
env_key
))
}
}
pub fn sanitize_function_name(name: &str) -> String {
let re = Regex::new(r"[^a-zA-Z0-9_-]").unwrap();
re.replace_all(name, "_").to_string()
}
pub fn is_valid_function_name(name: &str) -> bool {
let re = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();
re.is_match(name)
}
/// Extract the model name from a JSON object. Common with most providers to have this top level attribute.
pub fn get_model(data: &Value) -> String {
if let Some(model) = data.get("model") {
if let Some(model_str) = model.as_str() {
model_str.to_string()
} else {
"Unknown".to_string()
}
} else {
"Unknown".to_string()
}
}
/// Check if a file is actually an image by examining its magic bytes
fn is_image_file(path: &Path) -> bool {
if let Ok(mut file) = std::fs::File::open(path) {
let mut buffer = [0u8; 8]; // Large enough for most image magic numbers
if file.read(&mut buffer).is_ok() {
// Check magic numbers for common image formats
return match &buffer[0..4] {
// PNG: 89 50 4E 47
[0x89, 0x50, 0x4E, 0x47] => true,
// JPEG: FF D8 FF
[0xFF, 0xD8, 0xFF, _] => true,
// GIF: 47 49 46 38
[0x47, 0x49, 0x46, 0x38] => true,
_ => false,
};
}
}
false
}
/// Convert a local image file to base64 encoded ImageContent
pub fn load_image_file(path: &str) -> Result<ImageContent, ProviderError> {
let path = Path::new(path);
// Verify it's an image before proceeding
if !is_image_file(path) {
return Err(ProviderError::RequestFailed(
"File is not a valid image".to_string(),
));
}
// Read the file
let bytes = std::fs::read(path)
.map_err(|e| ProviderError::RequestFailed(format!("Failed to read image file: {}", e)))?;
// Detect mime type from extension
let mime_type = match path.extension().and_then(|e| e.to_str()) {
Some(ext) => match ext.to_lowercase().as_str() {
"png" => "image/png",
"jpg" | "jpeg" => "image/jpeg",
_ => {
return Err(ProviderError::RequestFailed(
"Unsupported image format".to_string(),
));
}
},
None => {
return Err(ProviderError::RequestFailed(
"Unknown image format".to_string(),
));
}
};
// Convert to base64
let data = base64::prelude::BASE64_STANDARD.encode(&bytes);
Ok(ImageContent {
mime_type: mime_type.to_string(),
data,
})
}
pub fn emit_debug_trace(
model_config: &ModelConfig,
payload: &Value,
response: &Value,
usage: &Usage,
) {
tracing::debug!(
model_config = %serde_json::to_string_pretty(model_config).unwrap_or_default(),
input = %serde_json::to_string_pretty(payload).unwrap_or_default(),
output = %serde_json::to_string_pretty(response).unwrap_or_default(),
input_tokens = ?usage.input_tokens.unwrap_or_default(),
output_tokens = ?usage.output_tokens.unwrap_or_default(),
total_tokens = ?usage.total_tokens.unwrap_or_default(),
);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sanitize_function_name() {
assert_eq!(sanitize_function_name("hello-world"), "hello-world");
assert_eq!(sanitize_function_name("hello world"), "hello_world");
assert_eq!(sanitize_function_name("hello@world"), "hello_world");
}
#[test]
fn test_is_valid_function_name() {
assert!(is_valid_function_name("hello-world"));
assert!(is_valid_function_name("hello_world"));
assert!(!is_valid_function_name("hello world"));
assert!(!is_valid_function_name("hello@world"));
}
}
@@ -1,32 +0,0 @@
use crate::{
providers::{create, errors::ProviderError, ProviderExtractResponse},
types::json_value_ffi::JsonValueFfi,
Message, ModelConfig,
};
/// Generates a structured output based on the provided schema,
/// system prompt and user messages.
#[uniffi::export(async_runtime = "tokio", default(request_id = None))]
pub async fn generate_structured_outputs(
provider_name: &str,
provider_config: JsonValueFfi,
system_prompt: &str,
messages: &[Message],
schema: JsonValueFfi,
request_id: Option<String>,
) -> Result<ProviderExtractResponse, ProviderError> {
// Use OpenAI models specifically for this task
let model_name = if provider_name == "databricks" {
"goose-gpt-4-1"
} else {
"gpt-4.1"
};
let model_cfg = ModelConfig::new(model_name.to_string()).with_temperature(Some(0.0));
let provider = create(provider_name, provider_config, model_cfg)?;
let resp = provider
.extract(system_prompt, messages, &schema, request_id.as_deref())
.await?;
Ok(resp)
}
-247
View File
@@ -1,247 +0,0 @@
// This file defines types for completion interfaces, including the request and response structures.
// Many of these are adapted based on the Goose Service API:
// https://docs.google.com/document/d/1r5vjSK3nBQU1cIRf0WKysDigqMlzzrzl_bxEE4msOiw/edit?tab=t.0
use std::collections::HashMap;
use thiserror::Error;
use serde::{Deserialize, Serialize};
use crate::types::json_value_ffi::JsonValueFfi;
use crate::{message::Message, providers::Usage};
use crate::{model::ModelConfig, providers::errors::ProviderError};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionRequest {
pub provider_name: String,
pub provider_config: serde_json::Value,
pub model_config: ModelConfig,
pub system_preamble: Option<String>,
pub system_prompt_override: Option<String>,
pub messages: Vec<Message>,
pub extensions: Vec<ExtensionConfig>,
pub request_id: Option<String>,
}
impl CompletionRequest {
pub fn new(
provider_name: String,
provider_config: serde_json::Value,
model_config: ModelConfig,
system_preamble: Option<String>,
system_prompt_override: Option<String>,
messages: Vec<Message>,
extensions: Vec<ExtensionConfig>,
) -> Self {
Self {
provider_name,
provider_config,
model_config,
system_prompt_override,
system_preamble,
messages,
extensions,
request_id: None,
}
}
pub fn with_request_id(mut self, request_id: String) -> Self {
self.request_id = Some(request_id);
self
}
}
#[allow(clippy::too_many_arguments)]
#[uniffi::export(default(system_preamble = None, system_prompt_override = None))]
pub fn create_completion_request(
provider_name: &str,
provider_config: JsonValueFfi,
model_config: ModelConfig,
system_preamble: Option<String>,
system_prompt_override: Option<String>,
messages: Vec<Message>,
extensions: Vec<ExtensionConfig>,
request_id: Option<String>,
) -> CompletionRequest {
let mut request = CompletionRequest::new(
provider_name.to_string(),
provider_config,
model_config,
system_preamble,
system_prompt_override,
messages,
extensions,
);
if let Some(req_id) = request_id {
request = request.with_request_id(req_id);
}
request
}
uniffi::custom_type!(CompletionRequest, String, {
lower: |tc: &CompletionRequest| {
serde_json::to_string(&tc).unwrap()
},
try_lift: |s: String| {
Ok(serde_json::from_str(&s).unwrap())
},
});
// https://mozilla.github.io/uniffi-rs/latest/proc_macro/errors.html
#[derive(Debug, Error, uniffi::Error)]
#[uniffi(flat_error)]
pub enum CompletionError {
#[error("failed to create provider: {0}")]
UnknownProvider(String),
#[error("provider error: {0}")]
Provider(#[from] ProviderError),
#[error("template rendering error: {0}")]
Template(#[from] minijinja::Error),
#[error("json serialization error: {0}")]
Json(#[from] serde_json::Error),
#[error("tool not found error: {0}")]
ToolNotFound(String),
}
#[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
pub struct CompletionResponse {
pub message: Message,
pub model: String,
pub usage: Usage,
pub runtime_metrics: RuntimeMetrics,
}
impl CompletionResponse {
pub fn new(
message: Message,
model: String,
usage: Usage,
runtime_metrics: RuntimeMetrics,
) -> Self {
Self {
message,
model,
usage,
runtime_metrics,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
pub struct RuntimeMetrics {
pub total_time_sec: f32,
pub total_time_sec_provider: f32,
pub tokens_per_second: Option<f64>,
}
impl RuntimeMetrics {
pub fn new(
total_time_sec: f32,
total_time_sec_provider: f32,
tokens_per_second: Option<f64>,
) -> Self {
Self {
total_time_sec,
total_time_sec_provider,
tokens_per_second,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Enum)]
pub enum ToolApprovalMode {
Auto,
Manual,
Smart,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
pub struct ToolConfig {
pub name: String,
pub description: String,
pub input_schema: JsonValueFfi,
pub approval_mode: ToolApprovalMode,
}
impl ToolConfig {
pub fn new(
name: &str,
description: &str,
input_schema: JsonValueFfi,
approval_mode: ToolApprovalMode,
) -> Self {
Self {
name: name.to_string(),
description: description.to_string(),
input_schema,
approval_mode,
}
}
/// Convert the tool config to a core tool
pub fn to_core_tool(&self, name: Option<&str>) -> super::core::Tool {
let tool_name = name.unwrap_or(&self.name);
super::core::Tool::new(
tool_name,
self.description.clone(),
self.input_schema.clone(),
)
}
}
#[uniffi::export]
pub fn create_tool_config(
name: &str,
description: &str,
input_schema: JsonValueFfi,
approval_mode: ToolApprovalMode,
) -> ToolConfig {
ToolConfig::new(name, description, input_schema, approval_mode)
}
// — Register the newtypes with UniFFI, converting via JSON strings —
#[derive(Debug, Clone, Serialize, Deserialize, uniffi::Record)]
pub struct ExtensionConfig {
name: String,
instructions: Option<String>,
tools: Vec<ToolConfig>,
}
impl ExtensionConfig {
pub fn new(name: String, instructions: Option<String>, tools: Vec<ToolConfig>) -> Self {
Self {
name,
instructions,
tools,
}
}
/// Convert the tools to core tools with the extension name as a prefix
pub fn get_prefixed_tools(&self) -> Vec<super::core::Tool> {
self.tools
.iter()
.map(|tool| {
let name = format!("{}__{}", self.name, tool.name);
tool.to_core_tool(Some(&name))
})
.collect()
}
/// Get a map of prefixed tool names to their approval modes
pub fn get_prefixed_tool_configs(&self) -> HashMap<String, ToolConfig> {
self.tools
.iter()
.map(|tool| {
let name = format!("{}__{}", self.name, tool.name);
(name, tool.clone())
})
.collect()
}
}
-131
View File
@@ -1,131 +0,0 @@
// This file defines core types that require serialization to
// construct payloads for LLM model providers and work with MCPs.
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Enum)]
#[serde(rename_all = "lowercase")]
pub enum Role {
User,
Assistant,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Enum)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum Content {
Text(TextContent),
Image(ImageContent),
}
impl Content {
pub fn text<S: Into<String>>(text: S) -> Self {
Content::Text(TextContent { text: text.into() })
}
pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
Content::Image(ImageContent {
data: data.into(),
mime_type: mime_type.into(),
})
}
/// Get the text content if this is a TextContent variant
pub fn as_text(&self) -> Option<&str> {
match self {
Content::Text(text) => Some(&text.text),
_ => None,
}
}
/// Get the image content if this is an ImageContent variant
pub fn as_image(&self) -> Option<(&str, &str)> {
match self {
Content::Image(image) => Some((&image.data, &image.mime_type)),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
#[serde(rename_all = "camelCase")]
pub struct TextContent {
pub text: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, uniffi::Record)]
#[serde(rename_all = "camelCase")]
pub struct ImageContent {
pub data: String,
pub mime_type: String,
}
/// A tool that can be used by a model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Tool {
/// The name of the tool
pub name: String,
/// A description of what the tool does
pub description: String,
/// A JSON Schema object defining the expected parameters for the tool
pub input_schema: serde_json::Value,
}
impl Tool {
/// Create a new tool with the given name and description
pub fn new<N, D>(name: N, description: D, input_schema: serde_json::Value) -> Self
where
N: Into<String>,
D: Into<String>,
{
Tool {
name: name.into(),
description: description.into(),
input_schema,
}
}
}
/// A tool call request that an extension can execute
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolCall {
/// The name of the tool to execute
pub name: String,
/// The parameters for the execution
pub arguments: serde_json::Value,
/// Whether the tool call needs approval before execution. Default is false.
pub needs_approval: bool,
}
impl ToolCall {
/// Create a new ToolUse with the given name and parameters
pub fn new<S: Into<String>>(name: S, arguments: serde_json::Value) -> Self {
Self {
name: name.into(),
arguments,
needs_approval: false,
}
}
/// Set needs_approval field
pub fn set_needs_approval(&mut self, flag: bool) {
self.needs_approval = flag;
}
}
#[non_exhaustive]
#[derive(Error, Debug, Clone, Deserialize, Serialize, PartialEq, uniffi::Error)]
pub enum ToolError {
#[error("Invalid parameters: {0}")]
InvalidParameters(String),
#[error("Execution failed: {0}")]
ExecutionError(String),
#[error("Schema error: {0}")]
SchemaError(String),
#[error("Tool not found: {0}")]
NotFound(String),
}
pub type ToolResult<T> = std::result::Result<T, ToolError>;
@@ -1,18 +0,0 @@
use serde_json::Value;
// `serde_json::Value` gets converted to a `String` to pass across the FFI.
// https://github.com/mozilla/uniffi-rs/blob/main/docs/manual/src/types/custom_types.md?plain=1
// https://github.com/mozilla/uniffi-rs/blob/c7f6caa3d1bf20f934346cefd8e82b5093f0dc6f/examples/custom-types/src/lib.rs#L63-L69
uniffi::custom_type!(Value, String, {
// Remote is required since 'Value' is from a different crate
remote,
lower: |obj| {
serde_json::to_string(&obj).unwrap()
},
try_lift: |val| {
Ok(serde_json::from_str(&val).unwrap() )
},
});
pub type JsonValueFfi = Value;
-3
View File
@@ -1,3 +0,0 @@
pub mod completion;
pub mod core;
pub mod json_value_ffi;
@@ -1,79 +0,0 @@
use anyhow::Result;
use dotenvy::dotenv;
use goose_llm::extractors::generate_session_name;
use goose_llm::message::Message;
use goose_llm::providers::errors::ProviderError;
fn should_run_test() -> Result<(), String> {
dotenv().ok();
if std::env::var("DATABRICKS_HOST").is_err() {
return Err("Missing DATABRICKS_HOST".to_string());
}
if std::env::var("DATABRICKS_TOKEN").is_err() {
return Err("Missing DATABRICKS_TOKEN".to_string());
}
Ok(())
}
async fn _generate_session_name(messages: &[Message]) -> Result<String, ProviderError> {
let provider_name = "databricks";
let provider_config = serde_json::json!({
"host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
"token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
});
generate_session_name(provider_name, provider_config, messages, None).await
}
#[tokio::test]
async fn test_generate_session_name_success() {
if should_run_test().is_err() {
println!("Skipping...");
return;
}
// Build a few messages with at least two user messages
let messages = vec![
Message::user().with_text("Hello, how are you?"),
Message::assistant().with_text("I'm fine, thanks!"),
Message::user().with_text("What's the weather in New York tomorrow?"),
];
let name = _generate_session_name(&messages)
.await
.expect("Failed to generate session name");
println!("Generated session name: {:?}", name);
// Should be non-empty and at most 4 words
let name = name.trim();
assert!(!name.is_empty(), "Name must not be empty");
let word_count = name.split_whitespace().count();
assert!(
word_count <= 4,
"Name must be 4 words or less, got {}: {}",
word_count,
name
)
}
#[tokio::test]
async fn test_generate_session_name_no_user() {
if should_run_test().is_err() {
println!("Skipping 'test_generate_session_name_no_user'. Databricks creds not set");
return;
}
// No user messages → expect ExecutionError
let messages = vec![
Message::assistant().with_text("System starting…"),
Message::assistant().with_text("All systems go."),
];
let err = _generate_session_name(&messages).await;
assert!(
matches!(err, Err(ProviderError::ExecutionError(_))),
"Expected ExecutionError when there are no user messages, got: {:?}",
err
);
}
-88
View File
@@ -1,88 +0,0 @@
use anyhow::Result;
use dotenvy::dotenv;
use goose_llm::extractors::generate_tooltip;
use goose_llm::message::{Message, MessageContent, ToolRequest};
use goose_llm::providers::errors::ProviderError;
use goose_llm::types::core::{Content, ToolCall};
use serde_json::json;
fn should_run_test() -> Result<(), String> {
dotenv().ok();
if std::env::var("DATABRICKS_HOST").is_err() {
return Err("Missing DATABRICKS_HOST".to_string());
}
if std::env::var("DATABRICKS_TOKEN").is_err() {
return Err("Missing DATABRICKS_TOKEN".to_string());
}
Ok(())
}
async fn _generate_tooltip(messages: &[Message]) -> Result<String, ProviderError> {
let provider_name = "databricks";
let provider_config = serde_json::json!({
"host": std::env::var("DATABRICKS_HOST").expect("Missing DATABRICKS_HOST"),
"token": std::env::var("DATABRICKS_TOKEN").expect("Missing DATABRICKS_TOKEN"),
});
generate_tooltip(provider_name, provider_config, messages, None).await
}
#[tokio::test]
async fn test_generate_tooltip_simple() {
if should_run_test().is_err() {
println!("Skipping...");
return;
}
// Two plain-text messages
let messages = vec![
Message::user().with_text("Hello, how are you?"),
Message::assistant().with_text("I'm fine, thanks! How can I help?"),
];
let tooltip = _generate_tooltip(&messages)
.await
.expect("Failed to generate tooltip");
println!("Generated tooltip: {:?}", tooltip);
assert!(!tooltip.trim().is_empty(), "Tooltip must not be empty");
assert!(
tooltip.len() < 100,
"Tooltip should be reasonably short (<100 chars)"
);
}
#[tokio::test]
async fn test_generate_tooltip_with_tools() {
if should_run_test().is_err() {
println!("Skipping...");
return;
}
// 1) Assistant message with a tool request
let mut tool_req_msg = Message::assistant();
let req = ToolRequest {
id: "1".to_string(),
tool_call: Ok(ToolCall::new("get_time", json!({"timezone": "UTC"}))).into(),
};
tool_req_msg.content.push(MessageContent::ToolReq(req));
// 2) User message with the tool response
let tool_resp_msg = Message::user().with_tool_response(
"1",
Ok(vec![Content::text("The current time is 12:00 UTC")]).into(),
);
let messages = vec![tool_req_msg, tool_resp_msg];
let tooltip = _generate_tooltip(&messages)
.await
.expect("Failed to generate tooltip");
println!("Generated tooltip (tools): {:?}", tooltip);
assert!(!tooltip.trim().is_empty(), "Tooltip must not be empty");
assert!(
tooltip.len() < 100,
"Tooltip should be reasonably short (<100 chars)"
);
}
@@ -1,382 +0,0 @@
use anyhow::Result;
use dotenvy::dotenv;
use goose_llm::message::{Message, MessageContent};
use goose_llm::providers::base::Provider;
use goose_llm::providers::errors::ProviderError;
use goose_llm::providers::{databricks, openai};
use goose_llm::types::core::{Content, Tool};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
#[derive(Debug, Clone, Copy)]
enum TestStatus {
Passed,
Skipped,
Failed,
}
impl std::fmt::Display for TestStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TestStatus::Passed => write!(f, ""),
TestStatus::Skipped => write!(f, "⏭️"),
TestStatus::Failed => write!(f, ""),
}
}
}
struct TestReport {
results: Mutex<HashMap<String, TestStatus>>,
}
impl TestReport {
fn new() -> Arc<Self> {
Arc::new(Self {
results: Mutex::new(HashMap::new()),
})
}
fn record_status(&self, provider: &str, status: TestStatus) {
let mut results = self.results.lock().unwrap();
results.insert(provider.to_string(), status);
}
fn record_pass(&self, provider: &str) {
self.record_status(provider, TestStatus::Passed);
}
fn record_skip(&self, provider: &str) {
self.record_status(provider, TestStatus::Skipped);
}
fn record_fail(&self, provider: &str) {
self.record_status(provider, TestStatus::Failed);
}
fn print_summary(&self) {
println!("\n============== Providers ==============");
let results = self.results.lock().unwrap();
let mut providers: Vec<_> = results.iter().collect();
providers.sort_by(|a, b| a.0.cmp(b.0));
for (provider, status) in providers {
println!("{} {}", status, provider);
}
println!("=======================================\n");
}
}
lazy_static::lazy_static! {
static ref TEST_REPORT: Arc<TestReport> = TestReport::new();
static ref ENV_LOCK: Mutex<()> = Mutex::new(());
}
/// Generic test harness for any Provider implementation
struct ProviderTester {
provider: Arc<dyn Provider>,
name: String,
}
impl ProviderTester {
fn new<T: Provider + Send + Sync + 'static>(provider: T, name: String) -> Self {
Self {
provider: Arc::new(provider),
name,
}
}
async fn test_basic_response(&self) -> Result<()> {
let message = Message::user().with_text("Just say hello!");
let response = self
.provider
.complete("You are a helpful assistant.", &[message], &[], None)
.await?;
// For a basic response, we expect a single text response
assert_eq!(
response.message.content.len(),
1,
"Expected single content item in response"
);
// Verify we got a text response
assert!(
matches!(response.message.content[0], MessageContent::Text(_)),
"Expected text response"
);
Ok(())
}
async fn test_tool_usage(&self) -> Result<()> {
let weather_tool = Tool::new(
"get_weather",
"Get the weather for a location",
serde_json::json!({
"type": "object",
"required": ["location"],
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA"
}
}
}),
);
let message = Message::user().with_text("What's the weather like in San Francisco?");
let response1 = self
.provider
.complete(
"You are a helpful weather assistant.",
&[message.clone()],
&[weather_tool.clone()],
None,
)
.await?;
println!("=== {}::reponse1 ===", self.name);
dbg!(&response1);
println!("===================");
// Verify we got a tool request
assert!(
response1
.message
.content
.iter()
.any(|content| matches!(content, MessageContent::ToolReq(_))),
"Expected tool request in response"
);
let id = &response1
.message
.content
.iter()
.filter_map(|message| message.as_tool_request())
.next_back()
.expect("got tool request")
.id;
let weather = Message::user().with_tool_response(
id,
Ok(vec![Content::text(
"
50°F°C
Precipitation: 0%
Humidity: 84%
Wind: 2 mph
Weather
Saturday 9:00 PM
Clear",
)])
.into(),
);
// Verify we construct a valid payload including the request/response pair for the next inference
let response2 = self
.provider
.complete(
"You are a helpful weather assistant.",
&[message, response1.message, weather],
&[weather_tool],
None,
)
.await?;
println!("=== {}::reponse2 ===", self.name);
dbg!(&response2);
println!("===================");
assert!(
response2
.message
.content
.iter()
.any(|content| matches!(content, MessageContent::Text(_))),
"Expected text for final response"
);
Ok(())
}
async fn test_context_length_exceeded_error(&self) -> Result<()> {
// Google Gemini has a really long context window
let large_message_content = if self.name.to_lowercase() == "google" {
"hello ".repeat(1_300_000)
} else {
"hello ".repeat(300_000)
};
let messages = vec![
Message::user().with_text("hi there. what is 2 + 2?"),
Message::assistant().with_text("hey! I think it's 4."),
Message::user().with_text(&large_message_content),
Message::assistant().with_text("heyy!!"),
// Messages before this mark should be truncated
Message::user().with_text("what's the meaning of life?"),
Message::assistant().with_text("the meaning of life is 42"),
Message::user().with_text(
"did I ask you what's 2+2 in this message history? just respond with 'yes' or 'no'",
),
];
// Test that we get ProviderError::ContextLengthExceeded when the context window is exceeded
let result = self
.provider
.complete("You are a helpful assistant.", &messages, &[], None)
.await;
// Print some debug info
println!("=== {}::context_length_exceeded_error ===", self.name);
dbg!(&result);
println!("===================");
// Ollama truncates by default even when the context window is exceeded
if self.name.to_lowercase() == "ollama" {
assert!(
result.is_ok(),
"Expected to succeed because of default truncation"
);
return Ok(());
}
assert!(
result.is_err(),
"Expected error when context window is exceeded"
);
assert!(
matches!(result.unwrap_err(), ProviderError::ContextLengthExceeded(_)),
"Expected error to be ContextLengthExceeded"
);
Ok(())
}
/// Run all provider tests
async fn run_test_suite(&self) -> Result<()> {
self.test_basic_response().await?;
self.test_tool_usage().await?;
self.test_context_length_exceeded_error().await?;
Ok(())
}
}
fn load_env() {
if let Ok(path) = dotenv() {
println!("Loaded environment from {:?}", path);
}
}
/// Helper function to run a provider test with proper error handling and reporting
async fn test_provider<F, T>(
name: &str,
required_vars: &[&str],
env_modifications: Option<HashMap<&str, Option<String>>>,
provider_fn: F,
) -> Result<()>
where
F: FnOnce() -> T,
T: Provider + Send + Sync + 'static,
{
// We start off as failed, so that if the process panics it is seen as a failure
TEST_REPORT.record_fail(name);
// Take exclusive access to environment modifications
let lock = ENV_LOCK.lock().unwrap();
load_env();
// Save current environment state for required vars and modified vars
let mut original_env = HashMap::new();
for &var in required_vars {
if let Ok(val) = std::env::var(var) {
original_env.insert(var, val);
}
}
if let Some(mods) = &env_modifications {
for &var in mods.keys() {
if let Ok(val) = std::env::var(var) {
original_env.insert(var, val);
}
}
}
// Apply any environment modifications
if let Some(mods) = &env_modifications {
for (&var, value) in mods.iter() {
match value {
Some(val) => std::env::set_var(var, val),
None => std::env::remove_var(var),
}
}
}
// Setup the provider
let missing_vars = required_vars.iter().any(|var| std::env::var(var).is_err());
if missing_vars {
println!("Skipping {} tests - credentials not configured", name);
TEST_REPORT.record_skip(name);
return Ok(());
}
let provider = provider_fn();
// Restore original environment
for (&var, value) in original_env.iter() {
std::env::set_var(var, value);
}
if let Some(mods) = env_modifications {
for &var in mods.keys() {
if !original_env.contains_key(var) {
std::env::remove_var(var);
}
}
}
std::mem::drop(lock);
let tester = ProviderTester::new(provider, name.to_string());
match tester.run_test_suite().await {
Ok(_) => {
TEST_REPORT.record_pass(name);
Ok(())
}
Err(e) => {
println!("{} test failed: {}", name, e);
TEST_REPORT.record_fail(name);
Err(e)
}
}
}
#[tokio::test]
async fn openai_complete() -> Result<()> {
test_provider(
"OpenAI",
&["OPENAI_API_KEY"],
None,
openai::OpenAiProvider::default,
)
.await
}
#[tokio::test]
async fn databricks_complete() -> Result<()> {
test_provider(
"Databricks",
&["DATABRICKS_HOST", "DATABRICKS_TOKEN"],
None,
databricks::DatabricksProvider::default,
)
.await
}
// Print the final test report
#[ctor::dtor]
fn print_test_report() {
TEST_REPORT.print_summary();
}
-195
View File
@@ -1,195 +0,0 @@
// tests/providers_extract.rs
use anyhow::Result;
use dotenvy::dotenv;
use goose_llm::message::Message;
use goose_llm::providers::base::Provider;
use goose_llm::providers::{databricks::DatabricksProvider, openai::OpenAiProvider};
use goose_llm::ModelConfig;
use serde_json::{json, Value};
use std::sync::Arc;
#[derive(Debug, PartialEq, Copy, Clone)]
enum ProviderType {
OpenAi,
Databricks,
}
impl ProviderType {
fn required_env(&self) -> &'static [&'static str] {
match self {
ProviderType::OpenAi => &["OPENAI_API_KEY"],
ProviderType::Databricks => &["DATABRICKS_HOST", "DATABRICKS_TOKEN"],
}
}
fn create_provider(&self, cfg: ModelConfig) -> Result<Arc<dyn Provider>> {
Ok(match self {
ProviderType::OpenAi => Arc::new(OpenAiProvider::from_env(cfg)),
ProviderType::Databricks => Arc::new(DatabricksProvider::from_env(cfg)),
})
}
}
fn check_required_env_vars(required: &[&str]) -> bool {
let missing: Vec<_> = required
.iter()
.filter(|&&v| std::env::var(v).is_err())
.cloned()
.collect();
if !missing.is_empty() {
println!("Skipping test; missing env vars: {:?}", missing);
false
} else {
true
}
}
// --- Shared inputs for "paper" task ---
const PAPER_SYSTEM: &str =
"You are an expert at structured data extraction. Extract the metadata of a research paper into JSON.";
const PAPER_TEXT: &str =
"Application of Quantum Algorithms in Interstellar Navigation: A New Frontier \
by Dr. Stella Voyager, Dr. Nova Star, Dr. Lyra Hunter. Abstract: This paper \
investigates the utilization of quantum algorithms to improve interstellar \
navigation systems. Keywords: Quantum algorithms, interstellar navigation, \
space-time anomalies, quantum superposition, quantum entanglement, space travel.";
fn paper_schema() -> Value {
json!({
"type": "object",
"properties": {
"title": { "type": "string" },
"authors": { "type": "array", "items": { "type": "string" } },
"abstract": { "type": "string" },
"keywords": { "type": "array", "items": { "type": "string" } }
},
"required": ["title","authors","abstract","keywords"],
"additionalProperties": false
})
}
// --- Shared inputs for "UI" task ---
const UI_SYSTEM: &str = "You are a UI generator AI. Convert the user input into a JSON-driven UI.";
const UI_TEXT: &str = "Make a User Profile Form";
fn ui_schema() -> Value {
json!({
"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
})
}
/// Generic runner for any extract task
async fn run_extract_test<F>(
provider_type: ProviderType,
model: &str,
system: &'static str,
user_text: &'static str,
schema: Value,
validate: F,
) -> Result<()>
where
F: Fn(&Value) -> bool,
{
dotenv().ok();
if !check_required_env_vars(provider_type.required_env()) {
return Ok(());
}
let cfg = ModelConfig::new(model.to_string()).with_temperature(Some(0.0));
let provider = provider_type.create_provider(cfg)?;
let msg = Message::user().with_text(user_text);
let resp = provider.extract(system, &[msg], &schema, None).await?;
println!("[{:?}] extract => {}", provider_type, resp.data);
assert!(
validate(&resp.data),
"{:?} failed validation on {}",
provider_type,
resp.data
);
Ok(())
}
/// Helper for the "paper" task
async fn run_extract_paper_test(provider: ProviderType, model: &str) -> Result<()> {
run_extract_test(
provider,
model,
PAPER_SYSTEM,
PAPER_TEXT,
paper_schema(),
|v| {
v.as_object()
.map(|o| {
["title", "authors", "abstract", "keywords"]
.iter()
.all(|k| o.contains_key(*k))
})
.unwrap_or(false)
},
)
.await
}
/// Helper for the "UI" task
async fn run_extract_ui_test(provider: ProviderType, model: &str) -> Result<()> {
run_extract_test(provider, model, UI_SYSTEM, UI_TEXT, ui_schema(), |v| {
v.as_object()
.and_then(|o| o.get("type").and_then(Value::as_str))
== Some("form")
})
.await
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn openai_extract_paper() -> Result<()> {
run_extract_paper_test(ProviderType::OpenAi, "gpt-4o").await
}
#[tokio::test]
async fn openai_extract_ui() -> Result<()> {
run_extract_ui_test(ProviderType::OpenAi, "gpt-4o").await
}
#[tokio::test]
async fn databricks_extract_paper() -> Result<()> {
run_extract_paper_test(ProviderType::Databricks, "goose-gpt-4-1").await
}
#[tokio::test]
async fn databricks_extract_ui() -> Result<()> {
run_extract_ui_test(ProviderType::Databricks, "goose-gpt-4-1").await
}
}
-3
View File
@@ -1,3 +0,0 @@
fn main() {
uniffi::uniffi_bindgen_main()
}