diff --git a/.env.sample b/.env.sample new file mode 100644 index 0000000..e5e2da8 --- /dev/null +++ b/.env.sample @@ -0,0 +1,3 @@ +VERTEX_PROJECT_ID= +VERTEX_LOCATION= +VERTEX_CREDENTIALS_PATH= \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index f5c2ea7..fd0d42d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,6 +29,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + [[package]] name = "alloc-no-stdlib" version = "2.0.4" @@ -65,6 +74,27 @@ version = "1.0.92" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74f37166d7d48a0284b99dd824694c26119c700b53bf0d1540cdb147dbdaaf13" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "async-channel" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81953c529336010edd6d8e358f886d9581267795c61b19475b71314bffa46d35" +dependencies = [ + "concurrent-queue", + "event-listener", + "futures-core", +] + [[package]] name = "async-compression" version = "0.4.17" @@ -136,10 +166,10 @@ dependencies = [ "axum-core", "bytes", "futures-util", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.5.0", "hyper-util", "itoa", "matchit", @@ -169,8 +199,8 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -191,8 +221,8 @@ dependencies = [ "axum-core", "bytes", "futures-util", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -212,8 +242,8 @@ dependencies = [ "axum", "bytes", "futures-core", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "matchit", "metrics", "metrics-exporter-prometheus", @@ -239,6 +269,18 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -329,10 +371,20 @@ dependencies = [ "iana-time-zone", "js-sys", "num-traits", + "serde", "wasm-bindgen", "windows-targets", ] +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -373,6 +425,31 @@ version = "0.8.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +[[package]] +name = "deadpool" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "421fe0f90f2ab22016f32a9881be5134fdd71c65298917084b0c7477cbc3856e" +dependencies = [ + "async-trait", + "deadpool-runtime", + "num_cpus", + "retain_mut", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + +[[package]] +name = "dotenv" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" + [[package]] name = "either" version = "1.13.0" @@ -404,6 +481,21 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "event-listener" +version = "2.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" + +[[package]] +name = "fastrand" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51093e27b0797c359783294ca4f0a911c270184cb10f85783b118614a1501be" +dependencies = [ + "instant", +] + [[package]] name = "fastrand" version = "2.1.1" @@ -498,6 +590,21 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +[[package]] +name = "futures-lite" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49a9d51ce47660b1e808d3c990b4709f2f415d928835a17dfd16991515c46bce" +dependencies = [ + "fastrand 1.9.0", + "futures-core", + "futures-io", + "memchr", + "parking", + "pin-project-lite", + "waker-fn", +] + [[package]] name = "futures-macro" version = "0.3.31" @@ -521,6 +628,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.31" @@ -539,6 +652,44 @@ dependencies = [ "slab", ] +[[package]] +name = "gcp_auth" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbf67f30198e045a039264c01fb44659ce82402d7771c50938beb41a5ac87733" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "home", + "http 1.1.0", + "http-body-util", + "hyper 1.5.0", + "hyper-rustls", + "hyper-util", + "ring", + "rustls-pemfile", + "serde", + "serde_json", + "thiserror", + "tokio", + "tracing", + "tracing-futures", + "url", +] + +[[package]] +name = "getrandom" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.9.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.2.15" @@ -547,7 +698,7 @@ checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", ] [[package]] @@ -562,6 +713,25 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" +[[package]] +name = "h2" +version = "0.3.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81fe527a889e1532da5c525686d96d4c2e74cdd345badf8dfef9f6b39dd5f5e8" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap 2.6.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.6" @@ -573,7 +743,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", + "http 1.1.0", "indexmap 2.6.0", "slab", "tokio", @@ -618,6 +788,26 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "home" +version = "0.5.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf" +dependencies = [ + "windows-sys 0.59.0", +] + +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.1.0" @@ -629,6 +819,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -636,7 +837,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.1.0", ] [[package]] @@ -647,8 +848,8 @@ checksum = "793429d76616a256bcb62c2a2ec2bed781c8307e797e2598c50010f2bee2544f" dependencies = [ "bytes", "futures-util", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "pin-project-lite", ] @@ -658,6 +859,27 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08a397c49fec283e3d6211adbe480be95aae5f304cfb923e9970e08956d5168a" +[[package]] +name = "http-types" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e9b187a72d63adbfba487f48095306ac823049cb504ee195541e91c7775f5ad" +dependencies = [ + "anyhow", + "async-channel", + "base64 0.13.1", + "futures-lite", + "http 0.2.12", + "infer", + "pin-project-lite", + "rand 0.7.3", + "serde", + "serde_json", + "serde_qs", + "serde_urlencoded", + "url", +] + [[package]] name = "httparse" version = "1.9.5" @@ -680,7 +902,9 @@ dependencies = [ "axum-extra", "axum-prometheus", "chrono", + "dotenv", "futures", + "gcp_auth", "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions", @@ -696,6 +920,31 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", + "wiremock", +] + +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.26", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", + "want", ] [[package]] @@ -707,9 +956,9 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.6", + "http 1.1.0", + "http-body 1.0.1", "httparse", "httpdate", "itoa", @@ -726,8 +975,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08afdbb5c31130e3034af566421053ab03787c640246a446327f550d11bcb333" dependencies = [ "futures-util", - "http", - "hyper", + "http 1.1.0", + "hyper 1.5.0", "hyper-util", "rustls", "rustls-native-certs", @@ -743,7 +992,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper", + "hyper 1.5.0", "hyper-util", "pin-project-lite", "tokio", @@ -758,7 +1007,7 @@ checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" dependencies = [ "bytes", "http-body-util", - "hyper", + "hyper 1.5.0", "hyper-util", "native-tls", "tokio", @@ -775,9 +1024,9 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http", - "http-body", - "hyper", + "http 1.1.0", + "http-body 1.0.1", + "hyper 1.5.0", "pin-project-lite", "socket2", "tokio", @@ -838,6 +1087,21 @@ dependencies = [ "hashbrown 0.15.0", ] +[[package]] +name = "infer" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64e9829a50b42bb782c1df523f78d332fe371b10c661e78b7a3c34b0198e9fac" + +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + [[package]] name = "ipnet" version = "2.10.1" @@ -949,9 +1213,9 @@ version = "0.15.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4f0c8427b39666bf970460908b213ec09b3b350f20c0c2eabcbba51704a08e6" dependencies = [ - "base64", + "base64 0.22.1", "http-body-util", - "hyper", + "hyper 1.5.0", "hyper-util", "indexmap 2.6.0", "ipnet", @@ -1011,7 +1275,7 @@ checksum = "80e04d1dcff3aae0704555fe5fee3bcfaf3d1fdf8a7e521d5b9d2b42acb52cec" dependencies = [ "hermit-abi", "libc", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", "windows-sys 0.52.0", ] @@ -1142,7 +1406,7 @@ checksum = "10a8a7f5f6ba7c1b286c2fbca0454eaba116f63bbe69ed250b642d36fbb04d80" dependencies = [ "async-trait", "bytes", - "http", + "http 1.1.0", "opentelemetry", "reqwest", ] @@ -1155,7 +1419,7 @@ checksum = "91cf61a1868dacc576bf2b2a1c3e9ab150af7272909e80085c3173384fe11f76" dependencies = [ "async-trait", "futures-core", - "http", + "http 1.1.0", "opentelemetry", "opentelemetry-http", "opentelemetry-proto", @@ -1200,7 +1464,7 @@ dependencies = [ "once_cell", "opentelemetry", "percent-encoding", - "rand", + "rand 0.8.5", "serde_json", "thiserror", "tokio", @@ -1213,6 +1477,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.3" @@ -1337,7 +1607,7 @@ dependencies = [ "libc", "once_cell", "raw-cpuid", - "wasi", + "wasi 0.11.0+wasi-snapshot-preview1", "web-sys", "winapi", ] @@ -1367,7 +1637,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fadfaed2cd7f389d0161bb73eeb07b7b78f8691047a6f3e73caaeae55310a4a6" dependencies = [ "bytes", - "rand", + "rand 0.8.5", "ring", "rustc-hash", "rustls", @@ -1400,6 +1670,19 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom 0.1.16", + "libc", + "rand_chacha 0.2.2", + "rand_core 0.5.1", + "rand_hc", +] + [[package]] name = "rand" version = "0.8.5" @@ -1407,8 +1690,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core 0.5.1", ] [[package]] @@ -1418,7 +1711,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom 0.1.16", ] [[package]] @@ -1427,7 +1729,16 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.15", +] + +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core 0.5.1", ] [[package]] @@ -1448,23 +1759,52 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" + [[package]] name = "reqwest" version = "0.12.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a77c62af46e79de0a562e1a9849205ffcb7fc1238876e9bd743357570e04046f" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-channel", "futures-core", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.6", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.5.0", "hyper-rustls", "hyper-tls", "hyper-util", @@ -1516,6 +1856,12 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "retain_mut" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4389f1d5789befaf6029ebd9f7dac4af7f7e3d61b69d4f30e2ac02b57e7712b0" + [[package]] name = "ring" version = "0.17.8" @@ -1524,7 +1870,7 @@ checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" dependencies = [ "cc", "cfg-if", - "getrandom", + "getrandom 0.2.15", "libc", "spin", "untrusted", @@ -1701,6 +2047,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_qs" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7715380eec75f029a4ef7de39a9200e0a63823176b759d055b613f5a87df6a6" +dependencies = [ + "percent-encoding", + "serde", + "thiserror", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1847,7 +2204,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0f2c9fc62d0beef6951ccffd757e241266a2c833136efbe35af6cd2567dca5b" dependencies = [ "cfg-if", - "fastrand", + "fastrand 2.1.1", "once_cell", "rustix", "windows-sys 0.59.0", @@ -1981,13 +2338,13 @@ dependencies = [ "async-stream", "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", - "h2", - "http", - "http-body", + "h2 0.4.6", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.5.0", "hyper-timeout", "hyper-util", "percent-encoding", @@ -2013,7 +2370,7 @@ dependencies = [ "indexmap 1.9.3", "pin-project", "pin-project-lite", - "rand", + "rand 0.8.5", "slab", "tokio", "tokio-util", @@ -2050,8 +2407,8 @@ checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5" dependencies = [ "bitflags", "bytes", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", "pin-project-lite", "tower-layer", @@ -2065,13 +2422,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8437150ab6bbc8c5f0f519e3d5ed4aa883a83dd4cdd3d1b21f9482936046cb97" dependencies = [ "async-compression", - "base64", + "base64 0.22.1", "bitflags", "bytes", "futures-core", "futures-util", - "http", - "http-body", + "http 1.1.0", + "http-body 1.0.1", "http-body-util", "http-range-header", "httpdate", @@ -2134,6 +2491,16 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -2213,6 +2580,7 @@ dependencies = [ "form_urlencoded", "idna", "percent-encoding", + "serde", ] [[package]] @@ -2221,7 +2589,7 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" dependencies = [ - "getrandom", + "getrandom 0.2.15", ] [[package]] @@ -2242,6 +2610,12 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "waker-fn" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "317211a0dc0ceedd78fb2ca9a44aed3d7b9b26f81870d485c07122b4350673b7" + [[package]] name = "want" version = "0.3.1" @@ -2251,6 +2625,12 @@ dependencies = [ "try-lock", ] +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -2490,6 +2870,28 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wiremock" +version = "0.5.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13a3a53eaf34f390dd30d7b1b078287dd05df2aa2e21a589ccb80f5c7253c2e9" +dependencies = [ + "assert-json-diff", + "async-trait", + "base64 0.21.7", + "deadpool", + "futures", + "futures-timer", + "http-types", + "hyper 0.14.32", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "zerocopy" version = "0.7.35" diff --git a/Cargo.toml b/Cargo.toml index 9784651..bd858c8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,3 +39,8 @@ axum-prometheus = "0.7.0" reqwest-streams = { version = "0.8.1", features = ["json"] } futures = "0.3.31" async-stream = "0.3.6" +gcp_auth = "0.12" +dotenv = "0.15" + +[dev-dependencies] +wiremock = "0.5" \ No newline at end of file diff --git a/README.md b/README.md index 445a7f6..f8857e8 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,89 @@ completion = client.chat.completions.create( ) ``` + +## ☁️ Google VertexAI Provider + +Hub now supports Google VertexAI, allowing you to leverage Google's powerful LLMs like Gemini through our unified API. + +### Configuration + +To use the Google VertexAI provider, you need to configure it in your `config.yaml` file. Here's a sample configuration: + +```yaml +providers: + - key: vertexai + type: vertexai + project_id: "" + location: "" + credentials_path: "/path/to/service-account.json" +models: + - key: gemini-pro + type: gemini-pro + provider: vertexai + +pipelines: + - name: default + type: chat + plugins: + - logging: + level: info + - tracing: + endpoint: "https://api.traceloop.com/v1/traces" + api_key: "your-traceloop-api-key" + - model-router: + models: + - gemini-pro +``` + + +### Example Usage + +Here's an example of how to use the VertexAI provider with a chat completion pipeline: + +```yaml +models: + - key: gemini-pro-vertex + type: gemini-pro + provider: vertexai + +pipelines: + - name: default + type: chat + plugins: + - logging: + level: info + - tracing: + endpoint: "https://api.traceloop.com/v1/traces" + api_key: "" + - model-router: + models: + - gemini-pro-vertex +``` + +In this example: + +- A model named gemini-pro-vertex is configured to use the vertexai provider. + +- The default pipeline is configured to route requests to this model. + +You can now make chat completion requests and embedding requests to any vertexai models: + +```python +client = OpenAI( + base_url="http://localhost:3000/api/v1", + api_key="dummy", # a dummy key +) + +# Creating a chat completion with Gemini +completion = client.chat.completions.create( + model="gemini-pro", # Use the Gemini model name + messages=[{"role": "user", "content": "Tell me a joke about opentelemetry"}], + max_tokens=1000, +) +``` + + ## 🌱 Contributing Whether big or small, we love contributions ❤️ Check out our guide to see how to [get started](https://traceloop.com/docs/hub/contributing/overview). diff --git a/src/config/models.rs b/src/config/models.rs index d992203..6b87241 100644 --- a/src/config/models.rs +++ b/src/config/models.rs @@ -13,6 +13,7 @@ pub struct Config { pub struct Provider { pub key: String, pub r#type: String, + #[serde(default)] pub api_key: String, #[serde(flatten)] pub params: HashMap, diff --git a/src/models/embeddings.rs b/src/models/embeddings.rs index 12d1581..6ab9eaa 100644 --- a/src/models/embeddings.rs +++ b/src/models/embeddings.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use super::usage::Usage; +pub use super::usage::Usage; #[derive(Deserialize, Serialize, Clone)] pub struct EmbeddingsRequest { diff --git a/src/providers/mod.rs b/src/providers/mod.rs index 6a6ece1..71b57b2 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -3,3 +3,4 @@ pub mod azure; pub mod openai; pub mod provider; pub mod registry; +pub mod vertexai; diff --git a/src/providers/registry.rs b/src/providers/registry.rs index 9e5663f..86fef80 100644 --- a/src/providers/registry.rs +++ b/src/providers/registry.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use crate::config::models::Provider as ProviderConfig; use crate::providers::{ anthropic::AnthropicProvider, azure::AzureProvider, openai::OpenAIProvider, provider::Provider, + vertexai::VertexAIProvider, }; pub struct ProviderRegistry { @@ -20,6 +21,7 @@ impl ProviderRegistry { "openai" => Arc::new(OpenAIProvider::new(config)), "anthropic" => Arc::new(AnthropicProvider::new(config)), "azure" => Arc::new(AzureProvider::new(config)), + "vertexai" => Arc::new(VertexAIProvider::new(config)), _ => continue, }; providers.insert(config.key.clone(), provider); diff --git a/src/providers/vertexai/mod.rs b/src/providers/vertexai/mod.rs new file mode 100644 index 0000000..f6307e4 --- /dev/null +++ b/src/providers/vertexai/mod.rs @@ -0,0 +1,7 @@ +mod models; +mod provider; + +#[cfg(test)] +mod tests; + +pub use provider::VertexAIProvider; diff --git a/src/providers/vertexai/models.rs b/src/providers/vertexai/models.rs new file mode 100644 index 0000000..93d26ea --- /dev/null +++ b/src/providers/vertexai/models.rs @@ -0,0 +1,397 @@ +use crate::config::constants::default_max_tokens; +use crate::models::chat::{ChatCompletion, ChatCompletionChoice}; +use crate::models::content::{ChatCompletionMessage, ChatMessageContent}; +use crate::models::embeddings::{ + Embeddings, EmbeddingsInput, EmbeddingsRequest, EmbeddingsResponse, +}; +use crate::models::streaming::{ChatCompletionChunk, Choice, ChoiceDelta}; +use crate::models::tool_calls::{ChatMessageToolCall, FunctionCall}; +use crate::models::usage::Usage; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIChatCompletionRequest { + #[serde(rename = "contents")] + pub contents: Vec, + #[serde(rename = "generation_config")] + pub generation_config: Option, + #[serde(rename = "tools")] + #[serde(skip_serializing_if = "Vec::is_empty")] + pub tools: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct GenerationConfig { + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(rename = "topK")] + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + #[serde(rename = "topP")] + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(rename = "candidateCount")] + #[serde(skip_serializing_if = "Option::is_none")] + pub candidate_count: Option, + #[serde(rename = "maxOutputTokens")] + #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct Content { + pub role: String, + pub parts: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct Part { + #[serde(skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(rename = "functionCall", skip_serializing_if = "Option::is_none")] + pub function_call: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIChatCompletionResponse { + pub candidates: Vec, + #[serde(rename = "usageMetadata")] + pub usage_metadata: Option, + #[serde(rename = "modelVersion")] + pub model_version: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct GenerateContentResponse { + pub content: Content, + #[serde(rename = "finishReason")] + pub finish_reason: String, + #[serde(rename = "safetyRatings")] + pub safety_ratings: Option>, + #[serde(rename = "avgLogprobs")] + pub avg_logprobs: Option, + #[serde(rename = "functionCall")] + pub function_call: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct UsageMetadata { + #[serde(rename = "promptTokenCount")] + pub prompt_token_count: i32, + #[serde(rename = "candidatesTokenCount")] + pub candidates_token_count: i32, + #[serde(rename = "totalTokenCount")] + pub total_token_count: i32, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct SafetyRating { + pub category: String, + pub probability: String, + #[serde(rename = "probabilityScore")] + pub probability_score: f32, + pub severity: String, + #[serde(rename = "severityScore")] + pub severity_score: f32, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIStreamChunk { + pub candidates: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct GenerateContentStreamResponse { + pub content: Option, + #[serde(rename = "finishReason")] + pub finish_reason: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIEmbeddingsRequest { + pub instances: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub parameters: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIEmbeddingInstance { + pub content: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_type: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct EmbeddingParameters { + #[serde(skip_serializing_if = "Option::is_none")] + pub auto_truncate: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIEmbeddingsResponse { + pub predictions: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIPrediction { + pub embeddings: VertexAIEmbeddingValues, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIEmbeddingValues { + pub values: Vec, + pub statistics: VertexAIEmbeddingStatistics, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIEmbeddingStatistics { + pub truncated: bool, + pub token_count: u32, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct Tool { + pub function_declarations: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct FunctionDeclaration { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub parameters: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAITool { + pub function_declarations: Vec, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexAIFunctionDeclaration { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub parameters: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub(crate) struct VertexFunctionCall { + pub name: String, + pub args: serde_json::Value, +} + +impl From for VertexAIChatCompletionRequest { + fn from(request: crate::models::chat::ChatCompletionRequest) -> Self { + let contents = request + .messages + .into_iter() + .map(|message| { + let text = match message.content { + Some(ChatMessageContent::String(text)) => text, + Some(ChatMessageContent::Array(parts)) => parts + .into_iter() + .map(|part| part.text) + .collect::>() + .join(" "), + None => String::new(), + }; + + Content { + role: match message.role.as_str() { + "user" => "user".to_string(), + "assistant" => "model".to_string(), + _ => "user".to_string(), + }, + parts: vec![Part { + text: Some(text), + function_call: None, + }], + } + }) + .collect(); + + let tools = if let Some(tools) = request.tools { + vec![VertexAITool { + function_declarations: tools + .into_iter() + .map(|tool| VertexAIFunctionDeclaration { + name: tool.function.name, + description: tool.function.description, + parameters: tool + .function + .parameters + .map(|p| serde_json::to_value(p).unwrap_or_default()), + }) + .collect(), + }] + } else { + Vec::new() + }; + + VertexAIChatCompletionRequest { + contents, + generation_config: Some(GenerationConfig { + temperature: request.temperature, + top_k: None, + top_p: request.top_p, + candidate_count: request.n.map(|n| n as i32), + max_output_tokens: request.max_tokens.or(Some(default_max_tokens())), + }), + tools, + } + } +} + +impl From for ChatCompletion { + fn from(response: VertexAIChatCompletionResponse) -> Self { + let choices = response + .candidates + .into_iter() + .enumerate() + .map(|(index, candidate)| { + let (content, tool_calls) = if let Some(part) = candidate.content.parts.first() { + match (&part.text, &part.function_call) { + (Some(text), None) => (ChatMessageContent::String(text.clone()), None), + (None, Some(func_call)) => ( + ChatMessageContent::String(String::new()), + Some(vec![ChatMessageToolCall { + id: uuid::Uuid::new_v4().to_string(), + function: FunctionCall { + name: func_call.name.clone(), + arguments: func_call.args.to_string(), + }, + r#type: "function".to_string(), + }]), + ), + _ => (ChatMessageContent::String(String::new()), None), + } + } else { + (ChatMessageContent::String(String::new()), None) + }; + + ChatCompletionChoice { + index: index as u32, + message: ChatCompletionMessage { + role: "assistant".to_string(), + content: Some(content), + name: None, + tool_calls, + }, + finish_reason: Some(candidate.finish_reason), + logprobs: None, + } + }) + .collect(); + + let usage = response + .usage_metadata + .map(|metadata| Usage { + prompt_tokens: metadata.prompt_token_count as u32, + completion_tokens: metadata.candidates_token_count as u32, + total_tokens: metadata.total_token_count as u32, + completion_tokens_details: None, + prompt_tokens_details: None, + }) + .unwrap_or_default(); + + ChatCompletion { + id: uuid::Uuid::new_v4().to_string(), + object: None, + created: None, + model: "gemini-pro".to_string(), + choices, + usage, + system_fingerprint: None, + } + } +} + +impl From for ChatCompletionChunk { + fn from(chunk: VertexAIStreamChunk) -> Self { + ChatCompletionChunk { + id: uuid::Uuid::new_v4().to_string(), + choices: chunk + .candidates + .into_iter() + .enumerate() + .map(|(index, candidate)| Choice { + delta: ChoiceDelta { + content: candidate + .content + .and_then(|c| c.parts.first().and_then(|p| p.text.clone())), + role: Some("assistant".to_string()), + tool_calls: None, + }, + finish_reason: candidate.finish_reason, + index: index as u32, + logprobs: None, + }) + .collect(), + created: chrono::Utc::now().timestamp(), + model: "gemini-pro".to_string(), + system_fingerprint: None, + service_tier: None, + usage: None, + } + } +} + +impl From for VertexAIEmbeddingsRequest { + fn from(request: EmbeddingsRequest) -> Self { + let instances = match request.input { + EmbeddingsInput::Single(text) => vec![VertexAIEmbeddingInstance { + content: text, + task_type: None, + }], + EmbeddingsInput::Multiple(texts) => texts + .into_iter() + .map(|text| VertexAIEmbeddingInstance { + content: text, + task_type: None, + }) + .collect(), + }; + + VertexAIEmbeddingsRequest { + instances, + parameters: Some(EmbeddingParameters { + auto_truncate: Some(true), + }), + } + } +} + +impl From for EmbeddingsResponse { + fn from(response: VertexAIEmbeddingsResponse) -> Self { + let token_count = response + .predictions + .first() + .map(|p| p.embeddings.statistics.token_count) + .unwrap_or(0); + + EmbeddingsResponse { + object: "list".to_string(), + data: response + .predictions + .into_iter() + .enumerate() + .map(|(index, prediction)| Embeddings { + object: "embedding".to_string(), + embedding: prediction.embeddings.values, + index, + }) + .collect(), + model: "text-embedding-004".to_string(), + usage: Usage { + prompt_tokens: token_count, + completion_tokens: 0, + total_tokens: token_count, + completion_tokens_details: None, + prompt_tokens_details: None, + }, + } + } +} diff --git a/src/providers/vertexai/provider.rs b/src/providers/vertexai/provider.rs new file mode 100644 index 0000000..b9f72c5 --- /dev/null +++ b/src/providers/vertexai/provider.rs @@ -0,0 +1,315 @@ +use super::models::{ + VertexAIChatCompletionRequest, VertexAIChatCompletionResponse, VertexAIEmbeddingsRequest, + VertexAIEmbeddingsResponse, VertexAIStreamChunk, +}; +use crate::config::constants::stream_buffer_size_bytes; +use crate::config::models::{ModelConfig, Provider as ProviderConfig}; +use crate::models::chat::{ChatCompletionRequest, ChatCompletionResponse}; +use crate::models::completion::{CompletionRequest, CompletionResponse}; +use crate::models::embeddings::{EmbeddingsRequest, EmbeddingsResponse}; +use crate::providers::provider::Provider; +use axum::async_trait; +use axum::http::StatusCode; +use futures::StreamExt; +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use reqwest::{ + header::{HeaderMap, HeaderValue}, + Client, +}; +use reqwest_streams::error::{StreamBodyError, StreamBodyKind}; +use reqwest_streams::JsonStreamResponse; +use serde_json::json; +use std::sync::Arc; +use tokio::sync::Mutex; + +pub struct VertexAIProvider { + config: ProviderConfig, + http_client: Client, + token_provider: Arc>>>, +} + +#[async_trait] +impl Provider for VertexAIProvider { + fn new(config: &ProviderConfig) -> Self { + Self { + config: config.clone(), + http_client: Client::new(), + token_provider: Arc::new(Mutex::new(None)), + } + } + + fn key(&self) -> String { + self.config.key.clone() + } + + fn r#type(&self) -> String { + "vertexai".to_string() + } + + async fn chat_completions( + &self, + payload: ChatCompletionRequest, + _model_config: &ModelConfig, + ) -> Result { + let model = payload.model.clone(); + let token = self.get_token().await?; + let request: VertexAIChatCompletionRequest = payload.clone().into(); + + let project_id = self + .config + .params + .get("project_id") + .ok_or(StatusCode::BAD_REQUEST)?; + let location = self + .config + .params + .get("location") + .ok_or(StatusCode::BAD_REQUEST)?; + + let endpoint = if payload.stream.unwrap_or(false) { + "streamGenerateContent" + } else { + "generateContent" + }; + + let url = format!( + "https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model}:{endpoint}", + location = location, + project_id = project_id, + model = model + ); + + let mut headers = HeaderMap::new(); + headers.insert( + "Authorization", + HeaderValue::from_str(&format!("Bearer {}", token)) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ); + headers.insert("Content-Type", HeaderValue::from_static("application/json")); + + let request_body = json!({ + "contents": request.contents, + "generation_config": request.generation_config, + "tools": request.tools, + }); + + let response = self + .http_client + .post(url) + .headers(headers) + .json(&request_body) + .send() + .await + .map_err(|e| { + eprintln!("VertexAI API request error: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let status = response.status(); + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + eprintln!("VertexAI API request error: {}", error_text); + return Err( + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) + ); + } + + if payload.stream.unwrap_or(false) { + let buffer_size = stream_buffer_size_bytes() * 10; + let stream = response + .json_array_stream::(buffer_size) + .map(|result| { + result.map(|chunk| chunk.into()).map_err(|e| { + StreamBodyError::new(StreamBodyKind::CodecError, Some(Box::new(e)), None) + }) + }); + + Ok(ChatCompletionResponse::Stream(Box::pin(stream))) + } else { + let response_text = response.text().await.map_err(|e| { + eprintln!("Failed to get response text: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let vertex_response: VertexAIChatCompletionResponse = + serde_json::from_str(&response_text).map_err(|e| { + eprintln!( + "Failed to parse response: {}. Response was: {}", + e, response_text + ); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(ChatCompletionResponse::NonStream(vertex_response.into())) + } + } + async fn completions( + &self, + _payload: CompletionRequest, + _model_config: &ModelConfig, + ) -> Result { + // this is not supported by gemini use chat_completion + unimplemented!() + } + + async fn embeddings( + &self, + payload: EmbeddingsRequest, + _model_config: &ModelConfig, + ) -> Result { + let model = payload.model.clone(); + + let token = self.get_token().await?; + + let request: VertexAIEmbeddingsRequest = payload.into(); + + let project_id = self + .config + .params + .get("project_id") + .ok_or(StatusCode::BAD_REQUEST)?; + let location = self + .config + .params + .get("location") + .ok_or(StatusCode::BAD_REQUEST)?; + + let url = format!( + "https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model}:predict", + location = location, + project_id = project_id, + model = model + ); + + let mut headers = HeaderMap::new(); + headers.insert( + "Authorization", + HeaderValue::from_str(&format!("Bearer {}", token)) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ); + headers.insert("Content-Type", HeaderValue::from_static("application/json")); + + let response = self + .http_client + .post(url) + .headers(headers) + .json(&request) + .send() + .await + .map_err(|e| { + eprintln!("VertexAI API embeddings request error: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let status = response.status(); + + if !status.is_success() { + let error_text = response.text().await.unwrap_or_default(); + eprintln!("VertexAI API embeddings error: {}", error_text); + return Err( + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR) + ); + } + + let response_bytes = response.bytes().await.map_err(|e| { + eprintln!("Failed to get response bytes: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let vertex_response: VertexAIEmbeddingsResponse = serde_json::from_slice(&response_bytes) + .map_err(|e| { + eprintln!( + "Failed to parse response JSON: {}. Response was: {}", + e, + String::from_utf8_lossy(&response_bytes) + ); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(vertex_response.into()) + } +} + +impl VertexAIProvider { + async fn get_token(&self) -> Result { + let provider = { + let guard = self.token_provider.lock().await; + + match guard.as_ref() { + Some(p) => p.clone(), + None => { + drop(guard); + let new_provider = self.ensure_token_provider().await?; + let mut guard = self.token_provider.lock().await; + *guard = Some(new_provider.clone()); + new_provider + } + } + }; + + let token = provider + .token(&["https://www.googleapis.com/auth/cloud-platform"]) + .await + .map_err(|e| { + eprintln!("Authentication error: {}", e); + StatusCode::UNAUTHORIZED + })?; + + Ok(token.as_str().to_string()) + } + + async fn ensure_token_provider(&self) -> Result, StatusCode> { + match self.config.params.get("credentials_path") { + Some(path) => { + let service_account = CustomServiceAccount::from_file(path).map_err(|e| { + eprintln!("Failed to create service account from file: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(Arc::new(service_account) as Arc) + } + None => { + let provider = gcp_auth::provider().await.map_err(|e| { + eprintln!("Failed to create default token provider: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + Ok(provider) + } + } + } + + #[cfg(test)] + pub(crate) fn construct_url( + &self, + model: &str, + endpoint: &str, + project_id: &str, + location: &str, + ) -> String { + format!( + "https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model}:{endpoint}", + location = location, + project_id = project_id, + model = model + ) + } + + #[cfg(test)] + pub(crate) fn get_endpoint(&self, request: &ChatCompletionRequest) -> String { + if request.stream.unwrap_or(false) { + "streamGenerateContent".to_string() + } else { + "generateContent".to_string() + } + } + + #[cfg(test)] + pub(crate) async fn create_headers(&self, token: &str) -> Result { + let mut headers = HeaderMap::new(); + headers.insert( + "Authorization", + HeaderValue::from_str(&format!("Bearer {}", token)) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?, + ); + headers.insert("Content-Type", HeaderValue::from_static("application/json")); + Ok(headers) + } +} diff --git a/src/providers/vertexai/tests.rs b/src/providers/vertexai/tests.rs new file mode 100644 index 0000000..816433e --- /dev/null +++ b/src/providers/vertexai/tests.rs @@ -0,0 +1,314 @@ +mod tests { + use crate::config::models::Provider as ProviderConfig; + use crate::models::chat::{ChatCompletion, ChatCompletionRequest}; + use crate::models::content::{ChatCompletionMessage, ChatMessageContent}; + use crate::models::embeddings::{EmbeddingsInput, EmbeddingsRequest}; + use crate::models::tool_definition::{FunctionDefinition, ToolDefinition}; + use crate::providers::provider::Provider; + use crate::providers::vertexai::models::{ + Content, GenerateContentResponse, Part, UsageMetadata, VertexAIChatCompletionRequest, + VertexAIChatCompletionResponse, VertexAIEmbeddingsRequest, VertexFunctionCall, + }; + use crate::providers::vertexai::provider::VertexAIProvider; + use serde_json::json; + use std::collections::HashMap; + + fn create_test_config() -> ProviderConfig { + let mut params = HashMap::new(); + params.insert("project_id".to_string(), "test-project".to_string()); + params.insert("location".to_string(), "us-central1".to_string()); + params.insert( + "credentials_path".to_string(), + "test-credentials.json".to_string(), + ); + + ProviderConfig { + key: "test-vertexai".to_string(), + r#type: "vertexai".to_string(), + api_key: "".to_string(), + params, + } + } + + fn create_test_chat_request() -> ChatCompletionRequest { + ChatCompletionRequest { + model: "gemini-pro".to_string(), + messages: vec![ChatCompletionMessage { + role: "user".to_string(), + content: Some(ChatMessageContent::String( + "What's the weather in London?".to_string(), + )), + name: None, + tool_calls: None, + }], + temperature: None, + top_p: None, + n: None, + stream: None, + stop: None, + max_tokens: None, + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + tool_choice: None, + tools: None, + user: None, + parallel_tool_calls: None, + } + } + + #[test] + fn test_chat_request_conversion() { + let chat_request = ChatCompletionRequest { + model: "gemini-pro".to_string(), + messages: vec![ + ChatCompletionMessage { + role: "user".to_string(), + content: Some(ChatMessageContent::String("Hello".to_string())), + name: None, + tool_calls: None, + }, + ChatCompletionMessage { + role: "assistant".to_string(), + content: Some(ChatMessageContent::String("Hi there!".to_string())), + name: None, + tool_calls: None, + }, + ], + temperature: Some(0.7), + top_p: Some(0.9), + n: Some(1), + stream: None, + max_tokens: Some(100), + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + user: None, + tool_choice: None, + tools: None, + stop: None, + parallel_tool_calls: None, + }; + + let vertex_request: VertexAIChatCompletionRequest = chat_request.into(); + + assert_eq!(vertex_request.contents.len(), 2); + assert_eq!(vertex_request.contents[0].role, "user"); + assert_eq!( + vertex_request.contents[0].parts[0].text, + Some("Hello".to_string()) + ); + assert_eq!(vertex_request.contents[1].role, "model"); + assert_eq!( + vertex_request.contents[1].parts[0].text, + Some("Hi there!".to_string()) + ); + + let gen_config = vertex_request.generation_config.unwrap(); + assert_eq!(gen_config.temperature, Some(0.7)); + assert_eq!(gen_config.top_p, Some(0.9)); + assert_eq!(gen_config.max_output_tokens, Some(100)); + } + + #[test] + fn test_chat_response_conversion() { + let vertex_response = VertexAIChatCompletionResponse { + candidates: vec![GenerateContentResponse { + content: Content { + role: "model".to_string(), + parts: vec![Part { + text: Some("Generated response".to_string()), + function_call: None, + }], + }, + finish_reason: "stop".to_string(), + safety_ratings: None, + avg_logprobs: None, + function_call: None, + }], + usage_metadata: Some(UsageMetadata { + prompt_token_count: 10, + candidates_token_count: 20, + total_token_count: 30, + }), + model_version: Some("v1".to_string()), + }; + + let chat_completion: ChatCompletion = vertex_response.into(); + + assert_eq!(chat_completion.choices.len(), 1); + assert_eq!(chat_completion.choices[0].index, 0); + assert_eq!( + chat_completion.choices[0].finish_reason, + Some("stop".to_string()) + ); + + if let Some(ChatMessageContent::String(content)) = + &chat_completion.choices[0].message.content + { + assert_eq!(content, "Generated response"); + } else { + panic!("Expected String content"); + } + } + + #[test] + fn test_embeddings_request_conversion() { + let embeddings_request = EmbeddingsRequest { + model: "textembedding-gecko".to_string(), + input: EmbeddingsInput::Multiple(vec![ + "First text".to_string(), + "Second text".to_string(), + ]), + user: None, + encoding_format: None, + }; + + let vertex_request: VertexAIEmbeddingsRequest = embeddings_request.into(); + + assert_eq!(vertex_request.instances.len(), 2); + assert_eq!(vertex_request.instances[0].content, "First text"); + assert_eq!(vertex_request.instances[1].content, "Second text"); + assert!(vertex_request.parameters.unwrap().auto_truncate.unwrap()); + } + + #[test] + fn test_streaming_endpoint_selection() { + let config = create_test_config(); + let provider = VertexAIProvider::new(&config); + + let mut request = create_test_chat_request(); + request.stream = Some(true); + + let endpoint = provider.get_endpoint(&request); + assert_eq!(endpoint, "streamGenerateContent"); + + request.stream = Some(false); + let endpoint = provider.get_endpoint(&request); + assert_eq!(endpoint, "generateContent"); + } + + #[test] + fn test_provider_initialization() { + let config = create_test_config(); + let provider = VertexAIProvider::new(&config); + + assert_eq!(provider.key(), "test-vertexai"); + assert_eq!(provider.r#type(), "vertexai"); + } + + #[test] + fn test_url_construction() { + let config = create_test_config(); + let provider = VertexAIProvider::new(&config); + let chat_url = provider.construct_url( + "gemini-pro", + "generateContent", + "test-project", + "us-central1", + ); + + assert_eq!( + chat_url, + "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent" + ); + } + + #[tokio::test] + async fn test_header_construction() { + let config = create_test_config(); + let provider = VertexAIProvider::new(&config); + + let result = provider.create_headers("test-token").await; + + assert!(result.is_ok()); + let headers = result.unwrap(); + assert_eq!(headers.get("Authorization").unwrap(), "Bearer test-token"); + assert_eq!(headers.get("Content-Type").unwrap(), "application/json"); + } + #[test] + fn test_embeddings_request_construction() { + let request = EmbeddingsRequest { + model: "textembedding-gecko".to_string(), + input: EmbeddingsInput::Single("test text".to_string()), + user: None, + encoding_format: None, + }; + + let vertex_request: VertexAIEmbeddingsRequest = request.into(); + assert_eq!(vertex_request.instances[0].content, "test text"); + assert!(vertex_request.parameters.unwrap().auto_truncate.unwrap()); + } + + #[test] + fn test_function_calling_request() { + let mut chat_request = create_test_chat_request(); + chat_request.tools = Some(vec![ToolDefinition { + function: FunctionDefinition { + name: "get_weather".to_string(), + description: Some("Get the current weather in a location".to_string()), + parameters: Some(HashMap::from([ + ("type".to_string(), json!("object")), + ( + "properties".to_string(), + json!({ + "location": { + "type": "string", + "description": "The city name" + } + }), + ), + ("required".to_string(), json!(["location"])), + ])), + strict: None, + }, + tool_type: "function".to_string(), + }]); + + let vertex_request: VertexAIChatCompletionRequest = chat_request.into(); + + assert!(!vertex_request.tools.is_empty()); + assert_eq!( + vertex_request.tools[0].function_declarations[0].name, + "get_weather" + ); + assert_eq!( + vertex_request.tools[0].function_declarations[0].description, + Some("Get the current weather in a location".to_string()) + ); + } + + #[test] + fn test_function_calling_response() { + let vertex_response = VertexAIChatCompletionResponse { + candidates: vec![GenerateContentResponse { + content: Content { + role: "model".to_string(), + parts: vec![Part { + text: None, + function_call: Some(VertexFunctionCall { + name: "get_weather".to_string(), + args: json!({"location": "London"}), + }), + }], + }, + finish_reason: "stop".to_string(), + safety_ratings: None, + avg_logprobs: None, + function_call: None, + }], + usage_metadata: None, + model_version: None, + }; + + let chat_completion: ChatCompletion = vertex_response.into(); + + let tool_calls = chat_completion.choices[0] + .message + .tool_calls + .as_ref() + .unwrap(); + assert_eq!(tool_calls[0].function.name, "get_weather"); + assert_eq!(tool_calls[0].function.arguments, r#"{"location":"London"}"#); + } +} diff --git a/tests/recordings/vertexai/chat_completion.json b/tests/recordings/vertexai/chat_completion.json new file mode 100644 index 0000000..719557c --- /dev/null +++ b/tests/recordings/vertexai/chat_completion.json @@ -0,0 +1,166 @@ +[ + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/gemini-pro:generateContent", + "headers": {}, + "body": { + "max_tokens": 100, + "messages": [ + { + "content": "What is the capital of France?", + "role": "user" + } + ], + "model": "gemini-pro", + "temperature": 0.699999988079071 + } + }, + "response": { + "status": 200, + "body": { + "choices": [ + { + "finish_reason": "STOP", + "index": 0, + "message": { + "content": "The capital of France is Paris.", + "role": "assistant" + } + } + ], + "id": "7277544b-4643-4d73-b162-eaa0fdef9514", + "model": "gemini-pro", + "system_fingerprint": null, + "usage": { + "completion_tokens": 7, + "prompt_tokens": 7, + "total_tokens": 14 + } + } + } + }, + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/gemini-pro:generateContent", + "headers": {}, + "body": { + "max_tokens": 100, + "messages": [ + { + "content": "What is the capital of France?", + "role": "user" + } + ], + "model": "gemini-pro", + "temperature": 0.699999988079071 + } + }, + "response": { + "status": 200, + "body": { + "choices": [ + { + "finish_reason": "STOP", + "index": 0, + "message": { + "content": "The capital of France is Paris.", + "role": "assistant" + } + } + ], + "id": "018a66d5-c0af-4805-8b10-9733fc78b60e", + "model": "gemini-pro", + "system_fingerprint": null, + "usage": { + "completion_tokens": 7, + "prompt_tokens": 7, + "total_tokens": 14 + } + } + } + }, + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/gemini-pro:generateContent", + "headers": {}, + "body": { + "max_tokens": 100, + "messages": [ + { + "content": "What is the capital of France?", + "role": "user" + } + ], + "model": "gemini-pro", + "temperature": 0.699999988079071 + } + }, + "response": { + "status": 200, + "body": { + "choices": [ + { + "finish_reason": "STOP", + "index": 0, + "message": { + "content": "The capital of France is Paris. \n\nWould you like to know anything else about France? I can provide you with information on its history, culture, geography, or anything else you might be interested in. \n", + "role": "assistant" + } + } + ], + "id": "208cb351-7fc7-4e98-bfb7-d4832a0d0420", + "model": "gemini-pro", + "system_fingerprint": null, + "usage": { + "completion_tokens": 44, + "prompt_tokens": 7, + "total_tokens": 51 + } + } + } + }, + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/gemini-pro:generateContent", + "headers": {}, + "body": { + "max_tokens": 100, + "messages": [ + { + "content": "What is the capital of France?", + "role": "user" + } + ], + "model": "gemini-pro", + "temperature": 0.699999988079071 + } + }, + "response": { + "status": 200, + "body": { + "choices": [ + { + "finish_reason": "STOP", + "index": 0, + "message": { + "content": "The capital of France is Paris.", + "role": "assistant" + } + } + ], + "id": "22fb7111-3d7f-42e5-b2f5-243c958d2f96", + "model": "gemini-pro", + "system_fingerprint": null, + "usage": { + "completion_tokens": 7, + "prompt_tokens": 7, + "total_tokens": 14 + } + } + } + } +] \ No newline at end of file diff --git a/tests/recordings/vertexai/chat_completion_with_tools.json b/tests/recordings/vertexai/chat_completion_with_tools.json new file mode 100644 index 0000000..d8c786d --- /dev/null +++ b/tests/recordings/vertexai/chat_completion_with_tools.json @@ -0,0 +1,74 @@ +[ + { + "request": { + "method": "POST", + "path": "gemini-pro:generateContent", + "headers": {}, + "body": { + "max_tokens": 100, + "messages": [ + { + "content": "What's the weather in San Francisco?", + "role": "user" + } + ], + "model": "gemini-pro", + "temperature": 0.699999988079071, + "tools": [ + { + "function": { + "description": "Get the current weather in a location", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "The location to get weather for", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + } + ] + } + }, + "response": { + "status": 200, + "body": { + "choices": [ + { + "finish_reason": "STOP", + "index": 0, + "message": { + "content": "", + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": "{\"location\":\"San Francisco\"}", + "name": "get_weather" + }, + "id": "b6ff6d69-0310-4215-9640-a55614975d9e", + "type": "function" + } + ] + } + } + ], + "id": "714a1e68-5dfa-4dd6-9861-342c7b3cc03d", + "model": "gemini-pro", + "system_fingerprint": null, + "usage": { + "completion_tokens": 6, + "prompt_tokens": 27, + "total_tokens": 33 + } + } + } + } +] \ No newline at end of file diff --git a/tests/recordings/vertexai/embeddings_test.json b/tests/recordings/vertexai/embeddings_test.json new file mode 100644 index 0000000..36b74f1 --- /dev/null +++ b/tests/recordings/vertexai/embeddings_test.json @@ -0,0 +1,1600 @@ +[ + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/textembedding-gecko:predict", + "headers": {}, + "body": { + "input": "Explore embeddings test.", + "model": "textembedding-gecko" + } + }, + "response": { + "status": 200, + "body": { + "data": [ + { + "embedding": [ + 0.024328794330358505, + -0.007927325554192066, + -0.03227777034044266, + 0.007646807469427586, + 0.042775582522153854, + 0.02543749287724495, + 0.039154794067144394, + -0.021833740174770355, + -0.02046365477144718, + 0.01565702073276043, + 0.03424655273556709, + -0.04123549163341522, + 0.009622358717024326, + -0.02369946427643299, + 0.039608296006917953, + -0.010863439179956911, + -0.015898792073130608, + -0.04297240450978279, + 0.034547194838523865, + -0.034599583595991135, + 0.02797384187579155, + 0.029799431562423703, + 0.05194639042019844, + -0.022435758262872696, + 0.02211863547563553, + 0.00012025993783026934, + 0.03602579981088638, + -0.05698813125491142, + -0.019475307315587997, + 0.044570449739694595, + -0.09668334573507308, + 0.02849353291094303, + -0.055953945964574814, + -0.004108349326997995, + -0.014041469432413578, + -0.048469699919223785, + -0.006792803760617971, + -0.001448061317205429, + -0.003982336726039648, + 0.04405730590224266, + 0.014645329676568508, + -0.0454498752951622, + -0.024309715256094933, + -0.0533146932721138, + 0.03288085386157036, + -0.01202381681650877, + 0.027210716158151627, + 0.03641979396343231, + -0.011927603743970394, + -0.04769068583846092, + -0.0003614430024754256, + -0.022049743682146072, + 0.04979113489389419, + -0.02043868973851204, + -0.007268783636391163, + -0.06732844561338425, + 0.03253961354494095, + -0.00685227382928133, + -0.02123664878308773, + 0.018453970551490784, + 0.01839176006615162, + -0.0330994576215744, + -0.03973471000790596, + 0.07108927518129349, + -0.0207611583173275, + -0.04231150075793266, + -0.023795245215296745, + 0.004485464189201593, + 0.04732951521873474, + -0.010377956554293633, + -0.0047367955558001995, + -0.03312591463327408, + 0.03884338587522507, + -0.030863812193274495, + -0.02829711139202118, + -0.055027831345796585, + 0.0043559931218624115, + 0.08080645650625229, + 0.03288064897060394, + 0.021732164546847343, + 0.007475654128938913, + 0.03323779255151749, + -0.04610563442111015, + -0.01843908242881298, + -0.03000701405107975, + 0.08127258718013763, + -0.03890974074602127, + -0.0066011808812618256, + -0.02887081913650036, + 0.05179654806852341, + 0.0017160448478534818, + -0.019606107845902443, + -0.0013793091056868434, + -0.08790423721075058, + 0.017591319978237152, + 0.032657135277986526, + -0.004456656984984875, + 0.015924790874123573, + -0.05703241750597954, + -0.0029068635776638985, + -0.005926548503339291, + 0.008181142620742321, + -0.0330340750515461, + 0.003892387729138136, + 0.029544739052653313, + -0.010662635788321497, + -0.004608854651451111, + 0.07462489604949951, + -0.029385121539235115, + 0.044994600117206573, + -0.02153143100440502, + -0.014342067763209345, + -0.05915575101971626, + -0.005340105388313532, + 0.05211803689599037, + -0.05388673022389412, + 0.008912160061299801, + 0.0711912214756012, + 0.05574534833431244, + -0.008361687883734703, + 0.027208322659134865, + -0.01887824945151806, + 0.06020164862275123, + -0.005311556626111269, + 0.014634489081799984, + 0.04607236012816429, + 0.04080691933631897, + 0.004521889612078667, + 0.043950404971838, + 0.06624354422092438, + -0.016513187438249588, + -0.06240186467766762, + 0.006469652988016605, + 0.03281101584434509, + 0.02143050730228424, + 0.05863906070590019, + 0.0460006482899189, + 0.03892882168292999, + 0.018882540985941887, + 0.021129202097654343, + 0.00149916869122535, + -0.006893137935549021, + -0.032483216375112534, + 0.04644506797194481, + -0.028843408450484276, + 0.004065012559294701, + -0.015335684642195702, + -0.0031936303712427616, + 0.01342240534722805, + -0.02130340412259102, + -0.05586519464850426, + 0.007106421049684286, + -0.08083653450012207, + -0.024053065106272697, + 0.019876092672348022, + -0.04103654995560646, + -0.016206173226237297, + 0.05481097847223282, + 0.014023763127624989, + 0.01893819309771061, + 0.04350786283612251, + -0.012183751910924911, + 0.015451154671609402, + -0.01502644270658493, + -0.013372528366744518, + 0.00358153460547328, + -0.008872230537235737, + 0.00620800256729126, + -0.060372766107320786, + -0.03989779204130173, + -0.02117086201906204, + 0.018981775268912315, + -0.03096655011177063, + -0.061312463134527206, + -0.030801720917224884, + -0.05938192084431648, + -0.0020013877656310797, + 0.003070233389735222, + -0.06251053512096405, + -0.032716233283281326, + -0.026364805176854137, + -0.04863995313644409, + 0.003774062963202596, + 0.01985598541796207, + 0.0631607174873352, + -0.02428035624325275, + 0.06702510267496109, + -0.016109831631183624, + -0.0029460121877491474, + 0.009330781176686289, + -0.03194137290120125, + -0.005103611387312412, + -0.03663048148155213, + -0.027177050709724423, + -0.05277325212955475, + -0.00699803838506341, + 0.005268934648483992, + 0.03343052417039871, + 0.005260680802166462, + -0.0036830201279371977, + 0.020199798047542572, + 0.07426322996616364, + 0.0030688082333654165, + 0.005452989600598812, + 0.018743395805358887, + -0.0224920604377985, + 0.10730841755867004, + -0.032293979078531265, + 0.0002350689173908904, + 0.0010789496591314671, + -0.03618810698390007, + 0.018345586955547333, + -0.04552971571683884, + 0.001177811180241406, + 0.00475353142246604, + 0.008068738505244255, + 0.0454033799469471, + -0.0291415061801672, + 0.008259382098913193, + -0.045275069773197174, + 0.014462602324783802, + -0.04790001362562179, + -0.00819101370871067, + -0.019015749916434288, + 0.012052169069647787, + 0.04507927969098091, + -0.03594772890210152, + 0.04199207201600075, + 0.021967053413391113, + -0.04701472818851471, + 0.042840376496315, + 0.08371294289827347, + 0.03654045611619949, + -0.03432486951351166, + 0.04463488608598709, + -0.03205568343400955, + -0.01638832688331604, + 0.01379961147904396, + 0.02304641157388687, + -0.01575622893869877, + -0.008729066699743271, + 0.06538544595241547, + 0.057345397770404816, + 0.00047449301928281784, + -0.031625501811504364, + -0.03249581903219223, + -0.00841601099818945, + 0.028201261535286903, + -0.05020326375961304, + 0.006538310553878546, + -0.008538326248526573, + -0.008140549063682556, + -0.005779254250228405, + 0.06331192702054977, + -0.08801457285881042, + 0.020517580211162567, + -0.033563870936632156, + -0.02176062949001789, + 0.00405531283468008, + 0.029410606250166893, + 0.0980067104101181, + 0.03568967059254646, + -0.017258282750844955, + 0.012557929381728172, + -0.011213832534849644, + -0.022580794990062714, + -0.0001237363467225805, + -0.10841508209705351, + -0.05358533188700676, + 0.008392318151891232, + -0.027041180059313778, + -0.06606398522853851, + 0.04253445193171501, + 0.04245651885867119, + -0.034624673426151276, + 0.06676711142063141, + 0.005221507977694273, + -0.0005807789275422692, + 0.03232831880450249, + -0.027065085247159004, + 0.04422459006309509, + 0.02007272094488144, + 0.023172909393906593, + -0.030005455017089844, + 0.007742203306406736, + 0.04723472520709038, + -0.04368368163704872, + -0.01076302770525217, + 0.007690794765949249, + -0.0006096992292441428, + -0.024208132177591324, + -0.009694481268525124, + 0.007394913118332624, + -0.04622989147901535, + -0.0026225668843835592, + 0.004615961108356714, + 0.012153132818639278, + 0.007988577708601952, + 0.02859876863658428, + 0.0005944292061030865, + 0.0198853500187397, + -0.02508392184972763, + -0.03584783151745796, + -0.02336323633790016, + 0.03417055681347847, + 0.018478572368621823, + -0.006408578250557184, + -0.05060793086886406, + 0.008679517544806004, + 0.01933043636381626, + -0.013685759156942368, + -0.02599271573126316, + -0.034445539116859436, + -0.02148745022714138, + 0.06461472809314728, + -0.023053385317325592, + 0.02325556054711342, + 0.010979730635881424, + -0.042654458433389664, + 0.07890661060810089, + -0.015566591173410416, + 0.07785733789205551, + -0.025808505713939667, + -0.006831650622189045, + -0.0035911358427256346, + -0.018313787877559665, + -0.05628412961959839, + 0.03278763219714165, + -0.014680576510727406, + 0.02322007156908512, + -0.009934262372553349, + -0.07628300786018372, + -0.0034670885652303696, + 0.011589798144996166, + 0.01053173840045929, + 0.05202061310410499, + -0.03773935511708259, + -0.02340039238333702, + -0.03212562948465347, + -0.057574111968278885, + -0.0077192275784909725, + 0.01510604377835989, + -0.035334497690200806, + -0.041306037455797195, + -0.00934187974780798, + 0.006484575569629669, + -0.01900274120271206, + -0.038125768303871155, + 0.07879000157117844, + 0.0003210725262761116, + 0.017772037535905838, + 0.05979960411787033, + -0.020214837044477463, + 0.0008718210156075656, + -0.001128549105487764, + -0.045359328389167786, + -0.008960873819887638, + 0.004953394643962383, + 0.03717462718486786, + -0.0027364802081137896, + -0.05447553098201752, + 0.024182872846722603, + 0.03637480363249779, + 0.0013033007271587849, + 0.02572665363550186, + 0.013465282507240772, + -0.012871243059635162, + 0.01009815651923418, + 0.006981448270380497, + 0.010489940643310549, + 0.0383855439722538, + -0.031269267201423645, + -0.0001566046994412318, + 8.68834467837587e-6, + 0.022414736449718475, + -0.014806135557591917, + -0.03377896547317505, + -0.005671503022313118, + 0.02651378884911537, + -0.035810623317956924, + -0.02546558901667595, + -0.029752327129244804, + 0.054470572620630264, + 0.03768720477819443, + -0.00477464497089386, + -0.02270832285284996, + 0.07969940453767776, + 0.037096939980983734, + 0.0060389903374016285, + 0.005754117388278246, + 0.0433335080742836, + 0.045446645468473434, + 0.08303327858448029, + 0.008431386202573776, + 0.017434200271964073, + -0.004610914271324873, + -0.023369504138827324, + -0.009266342036426067, + 0.02519710175693035, + -0.046084072440862656, + -0.01214672438800335, + 0.014546639285981657, + -0.08068262040615082, + -0.02034200355410576, + -0.010280669666826723, + -0.023736417293548584, + -0.04688943549990654, + -0.04601037874817848, + 0.021610168740153313, + 0.018097009509801865, + 0.02772311121225357, + 0.04196692258119583, + 0.030380651354789737, + -0.06747794896364212, + -0.0747714713215828, + -0.041327327489852905, + 0.043551649898290634, + -0.03888361528515816, + 0.03225799277424812, + 0.04595478996634483, + -0.02133955620229244, + 0.0019173286855220797, + 0.0059371874667704105, + -0.07239119708538055, + -0.045772816985845566, + -0.02472025156021118, + 0.03842487186193466, + -0.0498150922358036, + 0.015183750540018082, + 0.010547421872615814, + 0.004931791685521603, + 0.02025837264955044, + 0.01709268055856228, + 0.006867412012070417, + -0.012771272100508211, + -0.0314166285097599, + 0.01553863100707531, + -0.02006969042122364, + 0.003006607992574573, + -0.04408945515751839, + -0.011544267646968365, + -0.01973086409270763, + 0.03676016628742218, + 0.06280939280986786, + -0.046711549162864685, + -0.03145793080329895, + 0.01531656738370657, + -0.05163707584142685, + 0.05182184651494026, + -0.1592094600200653, + 0.027386890724301335, + -0.047682713717222214, + 0.0013601189712062478, + -0.06478971242904663, + -0.016342448070645332, + -0.015200386755168438, + -0.004773674067109823, + 0.08322356641292572, + 0.014807750470936298, + -0.015655919909477234, + -0.014719564467668532, + -0.006250931415706873, + -0.010224631056189535, + -0.1123192086815834, + -0.008011693134903908, + -0.03175446763634682, + 0.02245096117258072, + -0.0244610495865345, + 0.014179427176713943, + 0.04991292580962181, + 0.028357360512018204, + 0.003237012308090925, + 0.009162909351289272, + -0.004699984565377235, + -0.02256479114294052, + -0.012538294307887554, + -0.0866391658782959, + -0.003644892014563083, + 0.02131192944943905, + -0.03345496580004692, + -0.02353513427078724, + 0.01824165880680084, + 0.050371479243040085, + -0.01855982281267643, + -0.03273420408368111, + -0.037131503224372864, + 0.015415892004966736, + 0.020250506699085236, + 0.018454350531101227, + 0.05069306865334511, + -0.024432579055428505, + 0.016041111201047897, + 0.017883336171507835, + -0.05500530079007149, + 0.007667569909244776, + 0.036584850400686264, + 0.010960103943943976, + 0.05748739466071129, + 0.039034802466630936, + 0.022451212629675865, + 0.004143013618886471, + 0.004344919230788946, + 0.011412488296627998, + -0.03695736825466156, + 0.04837634041905403, + -0.024042043834924694, + 0.024671806022524834, + 0.011806590482592584, + 0.008030924946069717, + -0.0055960919708013535, + 0.015943756327033043, + 0.010229387320578098, + 0.051739875227212906, + -0.04000334441661835, + 0.010921076871454716, + -0.05255074426531792, + 0.011290939524769785, + 0.007231077644973993, + -0.010928617790341376, + -0.038236718624830246, + 0.08480462431907654, + 0.020943710580468174, + -0.0610310435295105, + -0.007050595711916685, + 0.04203200340270996, + -0.07865889370441437, + 0.03789631649851799, + 0.033277690410614014, + -0.028785821050405502, + -0.018483852967619896, + -0.011916191317141056, + 0.046082232147455215, + -0.0730227455496788, + -0.038391660898923874, + -0.021443212404847145, + -0.03609134256839752, + -0.012731975875794888, + 0.03956952318549156, + 0.05402698367834091, + -0.02773335762321949, + 0.01828986778855324, + -0.04690653830766678, + -0.005148555617779493, + 0.007142046000808477, + -0.01432266179472208, + -0.0015101222088560462, + 0.030206115916371346, + -0.08665747940540314, + 0.028638113290071487, + -0.01016109250485897, + -0.04504406079649925, + -0.002653571078553796, + 0.02551502361893654, + 0.004941532853990793, + 0.03649405390024185, + -0.04634816572070122, + -0.05233697593212128, + 0.02948020026087761, + 0.026992052793502808, + -0.010938028804957868, + -0.04512086883187294, + -0.026423104107379913, + 0.03986450657248497, + 0.0044611054472625256, + 0.023120151832699776, + 0.01017474103718996, + -0.02320389822125435, + -0.01825058087706566, + 0.10058684647083282, + 0.0012239363277330997, + -0.03815079852938652, + 0.023310359567403793, + -0.013754645362496376, + 0.06012607738375664, + 0.044174812734127045, + -0.0032098700758069754, + -0.020360436290502548, + -0.03826151043176651, + 0.011036238633096218, + -0.02455301582813263, + 0.019892487674951553, + 0.01371720340102911, + -0.008961028419435024, + 0.06835636496543884, + -0.013336542062461376, + 0.03928495943546295, + 0.06846529245376587, + 0.023292511701583862, + 0.015333494171500206, + 0.01914592646062374, + -0.06383553147315979, + 0.05568572133779526, + -0.045917607843875885, + -0.06129458546638489, + 0.062320876866579056, + 0.012246029451489449, + 0.004490720108151436, + -0.024292895570397377, + -0.02811212837696075, + -0.00669425493106246, + 0.01739179529249668, + -0.04005511477589607, + 0.09932595491409302, + -0.02828160487115383, + -0.016225267201662064, + 0.04903224855661392, + -0.011238664388656616, + -0.008266208693385124, + -0.0028693205676972866, + 0.04034503176808357, + -0.011773455888032911, + 0.027666568756103516, + 0.045300740748643875, + 0.03623325005173683, + -0.04606113955378533, + -0.00023302776389755309, + 0.010822237469255924, + 0.005396110936999321, + -0.017266767099499702, + -0.042832016944885254, + 0.000606586632784456, + 0.0053397659212350845, + 0.0031394879333674908, + -0.014797745272517204, + 0.04381722956895828, + 0.014831854961812496, + -0.028363173827528954, + -0.03383432328701019, + 0.031214935705065727, + 0.08290299773216248, + 0.02584749460220337, + 0.03909004107117653, + 0.0349324494600296, + -0.02809840254485607, + -0.05105957016348839, + -0.003338165581226349, + -0.044679734855890274, + -0.016897771507501602, + 0.04674547165632248, + -0.007903410121798515, + -0.06998921185731888, + -0.0004284674359951168, + 0.024417705833911896, + -0.025209281593561172, + 0.01556653156876564, + 0.0142686665058136, + 0.003608678700402379, + -0.03939566761255264, + -0.06773681193590164, + -0.006201360374689102, + -0.038590289652347565, + -0.04736294224858284, + 0.054148513823747635, + -0.04606981575489044, + -0.0216580331325531, + -0.02679860219359398, + -0.0387553870677948, + -0.06344524025917053, + 0.03479386121034622, + -0.03733435273170471, + -0.02730160392820835, + 0.033917054533958435, + -0.01180818397551775, + 0.017153600230813026, + -0.017416737973690033, + -0.05760662630200386, + 0.005898391827940941, + -0.06530486792325974, + -0.04456393793225288, + 0.013888120651245115, + -0.05623237416148186, + 0.012411084957420826, + 0.0311505813151598, + -0.05592577904462814, + 0.00470947427675128, + -0.012832939624786375, + -0.051747485995292664, + 0.08193240314722061, + 0.04231006279587746, + 0.060262277722358704, + -0.05159112438559532, + -0.0050200470723211765, + -0.028430450707674023, + 0.048510000109672546, + -0.03564736619591713, + 0.040754739195108414, + 0.028905360028147697, + 0.03899265825748443, + -0.0016340425936505198, + -0.008609502576291561, + 0.006150763016194105, + -0.004917451646178961, + -0.073238305747509, + 0.014989398419857023, + 0.039281148463487625, + 0.026978103443980217, + -0.01927020587027073, + 0.03466196730732918, + -0.026035647839307785, + 0.04892418906092644, + -0.021652784198522568, + -0.05802053213119507, + -0.007418050896376371, + 0.0004157012444920838, + -0.003113806713372469, + -0.04071711003780365, + -0.03618207946419716, + 0.02473904751241207, + 0.005009343847632408, + -0.003717175917699933, + 0.06939894706010818, + -0.03745422139763832, + -0.05634057521820069, + 0.020305491983890533, + 0.01624903455376625, + -0.0414578877389431, + -0.006078140810132027, + -0.00001147667171608191, + -0.03432951495051384, + 0.04420928284525871, + 0.03278480842709541, + 0.0019583101384341717, + 0.006309302523732185, + -0.041066475212574005, + 0.0045809498988091946, + 0.03785176947712898, + -0.030005238950252533, + 0.0525495707988739, + -0.024305732920765877, + 0.000991623615846038, + 0.026999732479453087, + -0.013087441213428974, + 0.0179459061473608, + 0.02947797439992428, + -0.0367741696536541, + 0.06801317632198334, + 0.0028329314664006233, + 0.02241651341319084, + -0.012403511442244051, + -0.06610151380300522, + -0.01063673384487629, + -0.0017124409787356851, + 0.009534181095659733, + 0.040957942605018616, + 0.030794164165854458, + -0.03456086665391922, + -0.04005061089992523, + -0.053644463419914246, + 0.011436519213020802, + -0.011531087569892406, + -0.020456653088331223, + -0.0051927464082837105, + 0.008158964104950428, + 0.02680368535220623, + 0.0762314721941948, + -0.022566702216863632, + 0.03298712149262428, + 0.0227197352796793, + 0.01228597667068243, + 0.0397644080221653, + -0.035839907824993134, + 0.04033581539988518, + -0.023282667621970177, + 0.0071014040149748325, + 0.060574986040592194, + 0.03509609028697014, + -0.023195691406726837, + 0.02417423389852047 + ], + "index": 0, + "object": "embedding" + } + ], + "model": "text-embedding-004", + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 7, + "total_tokens": 7 + } + } + } + }, + { + "request": { + "method": "POST", + "path": "/v1/projects/extended-legend-445620-u1/locations/us-central1/publishers/google/models/textembedding-gecko:predict", + "headers": {}, + "body": { + "input": "Explore embeddings test.", + "model": "textembedding-gecko" + } + }, + "response": { + "status": 200, + "body": { + "data": [ + { + "embedding": [ + 0.024328794330358505, + -0.007927325554192066, + -0.03227777034044266, + 0.007646807469427586, + 0.042775582522153854, + 0.02543749287724495, + 0.039154794067144394, + -0.021833740174770355, + -0.02046365477144718, + 0.01565702073276043, + 0.03424655273556709, + -0.04123549163341522, + 0.009622358717024326, + -0.02369946427643299, + 0.039608296006917953, + -0.010863439179956913, + -0.015898792073130608, + -0.04297240450978279, + 0.034547194838523865, + -0.034599583595991135, + 0.02797384187579155, + 0.029799431562423706, + 0.05194639042019844, + -0.022435758262872696, + 0.02211863547563553, + 0.00012025993783026934, + 0.03602579981088638, + -0.05698813125491142, + -0.019475307315587997, + 0.044570449739694595, + -0.09668334573507309, + 0.02849353291094303, + -0.055953945964574814, + -0.004108349326997995, + -0.014041469432413578, + -0.048469699919223785, + -0.006792803760617971, + -0.001448061317205429, + -0.003982336726039648, + 0.04405730590224266, + 0.014645329676568508, + -0.0454498752951622, + -0.024309715256094933, + -0.0533146932721138, + 0.03288085386157036, + -0.01202381681650877, + 0.027210716158151627, + 0.03641979396343231, + -0.011927603743970394, + -0.04769068583846092, + -0.0003614430024754256, + -0.022049743682146072, + 0.049791134893894196, + -0.02043868973851204, + -0.007268783636391163, + -0.06732844561338425, + 0.03253961354494095, + -0.00685227382928133, + -0.02123664878308773, + 0.018453970551490784, + 0.01839176006615162, + -0.0330994576215744, + -0.03973471000790596, + 0.07108927518129349, + -0.0207611583173275, + -0.04231150075793266, + -0.023795245215296745, + 0.004485464189201593, + 0.04732951521873474, + -0.010377956554293633, + -0.0047367955558001995, + -0.03312591463327408, + 0.03884338587522507, + -0.030863812193274498, + -0.02829711139202118, + -0.055027831345796585, + 0.0043559931218624115, + 0.08080645650625229, + 0.03288064897060394, + 0.021732164546847343, + 0.007475654128938913, + 0.03323779255151749, + -0.04610563442111015, + -0.01843908242881298, + -0.03000701405107975, + 0.08127258718013763, + -0.03890974074602127, + -0.0066011808812618256, + -0.02887081913650036, + 0.05179654806852341, + 0.0017160448478534818, + -0.019606107845902443, + -0.0013793091056868434, + -0.08790423721075058, + 0.017591319978237152, + 0.032657135277986526, + -0.004456656984984875, + 0.015924790874123573, + -0.05703241750597954, + -0.0029068635776638985, + -0.005926548503339291, + 0.008181142620742321, + -0.0330340750515461, + 0.003892387729138136, + 0.029544739052653313, + -0.010662635788321495, + -0.004608854651451111, + 0.07462489604949951, + -0.029385121539235115, + 0.044994600117206573, + -0.02153143100440502, + -0.014342067763209343, + -0.05915575101971626, + -0.005340105388313532, + 0.05211803689599037, + -0.05388673022389412, + 0.008912160061299801, + 0.0711912214756012, + 0.05574534833431244, + -0.008361687883734703, + 0.027208322659134865, + -0.01887824945151806, + 0.060201648622751236, + -0.005311556626111269, + 0.014634489081799984, + 0.04607236012816429, + 0.04080691933631897, + 0.004521889612078667, + 0.043950404971838, + 0.06624354422092438, + -0.016513187438249588, + -0.06240186467766762, + 0.006469652988016605, + 0.03281101584434509, + 0.02143050730228424, + 0.05863906070590019, + 0.0460006482899189, + 0.03892882168292999, + 0.018882540985941887, + 0.021129202097654343, + 0.00149916869122535, + -0.006893137935549021, + -0.032483216375112534, + 0.04644506797194481, + -0.028843408450484276, + 0.004065012559294701, + -0.015335684642195702, + -0.0031936303712427616, + 0.01342240534722805, + -0.02130340412259102, + -0.05586519464850426, + 0.007106421049684286, + -0.08083653450012207, + -0.024053065106272697, + 0.019876092672348022, + -0.04103654995560646, + -0.016206173226237297, + 0.05481097847223282, + 0.014023763127624989, + 0.01893819309771061, + 0.04350786283612251, + -0.012183751910924911, + 0.015451154671609402, + -0.01502644270658493, + -0.013372528366744518, + 0.00358153460547328, + -0.008872230537235737, + 0.00620800256729126, + -0.060372766107320786, + -0.03989779204130173, + -0.021170862019062042, + 0.018981775268912315, + -0.03096655011177063, + -0.061312463134527206, + -0.030801720917224884, + -0.05938192084431648, + -0.0020013877656310797, + 0.003070233389735222, + -0.06251053512096405, + -0.032716233283281326, + -0.026364805176854134, + -0.04863995313644409, + 0.0037740629632025957, + 0.019855985417962074, + 0.0631607174873352, + -0.024280356243252754, + 0.06702510267496109, + -0.016109831631183624, + -0.0029460121877491474, + 0.009330781176686287, + -0.03194137290120125, + -0.005103611387312412, + -0.036630481481552124, + -0.027177050709724426, + -0.05277325212955475, + -0.00699803838506341, + 0.005268934648483992, + 0.03343052417039871, + 0.005260680802166462, + -0.0036830201279371977, + 0.020199798047542572, + 0.07426322996616364, + 0.0030688082333654165, + 0.005452989600598812, + 0.018743395805358887, + -0.0224920604377985, + 0.10730841755867004, + -0.032293979078531265, + 0.00023506891739089042, + 0.0010789496591314673, + -0.03618810698390007, + 0.018345586955547333, + -0.04552971571683884, + 0.001177811180241406, + 0.00475353142246604, + 0.008068738505244255, + 0.0454033799469471, + -0.029141506180167198, + 0.008259382098913193, + -0.045275069773197174, + 0.014462602324783802, + -0.047900013625621796, + -0.00819101370871067, + -0.019015749916434288, + 0.012052169069647789, + 0.04507927969098091, + -0.03594772890210152, + 0.04199207201600075, + 0.021967053413391113, + -0.04701472818851471, + 0.042840376496315, + 0.08371294289827347, + 0.03654045611619949, + -0.03432486951351166, + 0.04463488608598709, + -0.03205568343400955, + -0.01638832688331604, + 0.01379961147904396, + 0.02304641157388687, + -0.01575622893869877, + -0.008729066699743271, + 0.06538544595241547, + 0.057345397770404816, + 0.00047449301928281784, + -0.031625501811504364, + -0.03249581903219223, + -0.00841601099818945, + 0.028201261535286903, + -0.05020326375961304, + 0.006538310553878546, + -0.008538326248526573, + -0.008140549063682556, + -0.005779254250228405, + 0.06331192702054977, + -0.08801457285881042, + 0.020517580211162567, + -0.033563870936632156, + -0.02176062949001789, + 0.00405531283468008, + 0.029410606250166893, + 0.0980067104101181, + 0.03568967059254646, + -0.017258282750844955, + 0.012557929381728172, + -0.011213832534849644, + -0.022580794990062714, + -0.0001237363467225805, + -0.10841508209705353, + -0.05358533188700676, + 0.008392318151891232, + -0.027041180059313774, + -0.06606398522853851, + 0.04253445193171501, + 0.04245651885867119, + -0.034624673426151276, + 0.06676711142063141, + 0.005221507977694273, + -0.0005807789275422692, + 0.03232831880450249, + -0.027065085247159004, + 0.04422459006309509, + 0.02007272094488144, + 0.023172909393906593, + -0.030005455017089844, + 0.007742203306406736, + 0.04723472520709038, + -0.04368368163704872, + -0.01076302770525217, + 0.007690794765949249, + -0.0006096992292441428, + -0.024208132177591324, + -0.009694481268525124, + 0.0073949131183326244, + -0.04622989147901535, + -0.0026225668843835592, + 0.004615961108356714, + 0.012153132818639278, + 0.007988577708601952, + 0.028598768636584282, + 0.0005944292061030865, + 0.0198853500187397, + -0.02508392184972763, + -0.03584783151745796, + -0.023363236337900162, + 0.03417055681347847, + 0.018478572368621826, + -0.006408578250557184, + -0.05060793086886406, + 0.008679517544806004, + 0.01933043636381626, + -0.013685759156942368, + -0.02599271573126316, + -0.034445539116859436, + -0.02148745022714138, + 0.06461472809314728, + -0.023053385317325592, + 0.02325556054711342, + 0.010979730635881424, + -0.042654458433389664, + 0.07890661060810089, + -0.015566591173410416, + 0.07785733789205551, + -0.025808505713939667, + -0.006831650622189045, + -0.0035911358427256346, + -0.018313787877559662, + -0.05628412961959839, + 0.03278763219714165, + -0.014680576510727406, + 0.02322007156908512, + -0.009934262372553349, + -0.07628300786018372, + -0.0034670885652303696, + 0.011589798144996166, + 0.01053173840045929, + 0.052020613104104996, + -0.037739355117082596, + -0.02340039238333702, + -0.03212562948465347, + -0.057574111968278885, + -0.0077192275784909725, + 0.01510604377835989, + -0.035334497690200806, + -0.041306037455797195, + -0.00934187974780798, + 0.006484575569629669, + -0.01900274120271206, + -0.038125768303871155, + 0.07879000157117844, + 0.0003210725262761116, + 0.017772037535905838, + 0.05979960411787033, + -0.020214837044477463, + 0.0008718210156075656, + -0.0011285491054877639, + -0.045359328389167786, + -0.008960873819887638, + 0.004953394643962383, + 0.03717462718486786, + -0.0027364802081137896, + -0.05447553098201752, + 0.024182872846722603, + 0.03637480363249779, + 0.0013033007271587849, + 0.02572665363550186, + 0.013465282507240772, + -0.012871243059635162, + 0.01009815651923418, + 0.006981448270380497, + 0.010489940643310547, + 0.0383855439722538, + -0.031269267201423645, + -0.0001566046994412318, + 8.68834467837587e-6, + 0.022414736449718475, + -0.014806135557591915, + -0.03377896547317505, + -0.005671503022313118, + 0.02651378884911537, + -0.035810623317956924, + -0.02546558901667595, + -0.029752327129244804, + 0.054470572620630264, + 0.03768720477819443, + -0.00477464497089386, + -0.02270832285284996, + 0.07969940453767776, + 0.037096939980983734, + 0.0060389903374016285, + 0.005754117388278246, + 0.0433335080742836, + 0.045446645468473434, + 0.08303327858448029, + 0.008431386202573776, + 0.017434200271964073, + -0.004610914271324873, + -0.023369504138827324, + -0.009266342036426067, + 0.02519710175693035, + -0.046084072440862656, + -0.01214672438800335, + 0.014546639285981655, + -0.08068262040615082, + -0.02034200355410576, + -0.010280669666826725, + -0.023736417293548584, + -0.04688943549990654, + -0.04601037874817848, + 0.021610168740153313, + 0.018097009509801865, + 0.02772311121225357, + 0.04196692258119583, + 0.030380651354789734, + -0.06747794896364212, + -0.0747714713215828, + -0.041327327489852905, + 0.043551649898290634, + -0.03888361528515816, + 0.03225799277424812, + 0.04595478996634483, + -0.021339556202292442, + 0.0019173286855220795, + 0.0059371874667704105, + -0.07239119708538055, + -0.045772816985845566, + -0.02472025156021118, + 0.03842487186193466, + -0.049815092235803604, + 0.015183750540018082, + 0.010547421872615814, + 0.004931791685521603, + 0.020258372649550438, + 0.01709268055856228, + 0.006867412012070417, + -0.012771272100508213, + -0.0314166285097599, + 0.01553863100707531, + -0.02006969042122364, + 0.0030066079925745726, + -0.04408945515751839, + -0.011544267646968365, + -0.019730864092707634, + 0.03676016628742218, + 0.06280939280986786, + -0.046711549162864685, + -0.03145793080329895, + 0.01531656738370657, + -0.05163707584142685, + 0.05182184651494026, + -0.1592094600200653, + 0.027386890724301338, + -0.047682713717222214, + 0.0013601189712062478, + -0.06478971242904663, + -0.016342448070645332, + -0.015200386755168438, + -0.004773674067109823, + 0.08322356641292572, + 0.014807750470936298, + -0.015655919909477234, + -0.014719564467668533, + -0.006250931415706873, + -0.010224631056189537, + -0.1123192086815834, + -0.008011693134903908, + -0.03175446763634682, + 0.02245096117258072, + -0.0244610495865345, + 0.014179427176713943, + 0.04991292580962181, + 0.028357360512018204, + 0.003237012308090925, + 0.009162909351289272, + -0.004699984565377235, + -0.02256479114294052, + -0.012538294307887554, + -0.0866391658782959, + -0.0036448920145630836, + 0.02131192944943905, + -0.03345496580004692, + -0.02353513427078724, + 0.018241658806800842, + 0.050371479243040085, + -0.01855982281267643, + -0.03273420408368111, + -0.037131503224372864, + 0.015415892004966736, + 0.020250506699085236, + 0.018454350531101227, + 0.05069306865334511, + -0.024432579055428505, + 0.016041111201047897, + 0.017883336171507835, + -0.05500530079007149, + 0.007667569909244776, + 0.036584850400686264, + 0.010960103943943977, + 0.05748739466071129, + 0.039034802466630936, + 0.022451212629675865, + 0.004143013618886471, + 0.004344919230788946, + 0.011412488296627998, + -0.03695736825466156, + 0.04837634041905403, + -0.024042043834924698, + 0.024671806022524834, + 0.011806590482592583, + 0.008030924946069717, + -0.0055960919708013535, + 0.015943756327033043, + 0.010229387320578098, + 0.051739875227212906, + -0.04000334441661835, + 0.010921076871454716, + -0.05255074426531792, + 0.011290939524769783, + 0.007231077644973993, + -0.010928617790341377, + -0.038236718624830246, + 0.08480462431907654, + 0.020943710580468178, + -0.0610310435295105, + -0.007050595711916685, + 0.04203200340270996, + -0.07865889370441437, + 0.03789631649851799, + 0.033277690410614014, + -0.028785821050405502, + -0.018483852967619896, + -0.011916191317141056, + 0.046082232147455215, + -0.0730227455496788, + -0.038391660898923874, + -0.021443212404847145, + -0.03609134256839752, + -0.012731975875794888, + 0.03956952318549156, + 0.05402698367834091, + -0.02773335762321949, + 0.018289867788553238, + -0.04690653830766678, + -0.005148555617779493, + 0.007142046000808477, + -0.01432266179472208, + -0.0015101222088560462, + 0.030206115916371346, + -0.08665747940540314, + 0.028638113290071487, + -0.01016109250485897, + -0.04504406079649925, + -0.002653571078553796, + 0.02551502361893654, + 0.004941532853990793, + 0.03649405390024185, + -0.04634816572070122, + -0.05233697593212128, + 0.02948020026087761, + 0.026992052793502808, + -0.010938028804957867, + -0.04512086883187294, + -0.026423104107379913, + 0.03986450657248497, + 0.0044611054472625256, + 0.023120151832699776, + 0.01017474103718996, + -0.02320389822125435, + -0.01825058087706566, + 0.10058684647083282, + 0.0012239363277330995, + -0.03815079852938652, + 0.023310359567403793, + -0.013754645362496376, + 0.06012607738375664, + 0.044174812734127045, + -0.0032098700758069754, + -0.020360436290502548, + -0.03826151043176651, + 0.011036238633096218, + -0.02455301582813263, + 0.019892487674951553, + 0.01371720340102911, + -0.008961028419435024, + 0.06835636496543884, + -0.013336542062461376, + 0.03928495943546295, + 0.06846529245376587, + 0.023292511701583862, + 0.015333494171500206, + 0.01914592646062374, + -0.06383553147315979, + 0.05568572133779526, + -0.045917607843875885, + -0.06129458546638489, + 0.062320876866579056, + 0.012246029451489449, + 0.004490720108151436, + -0.024292895570397377, + -0.028112128376960754, + -0.00669425493106246, + 0.01739179529249668, + -0.04005511477589607, + 0.09932595491409302, + -0.02828160487115383, + -0.016225267201662064, + 0.04903224855661392, + -0.011238664388656616, + -0.008266208693385124, + -0.0028693205676972866, + 0.04034503176808357, + -0.011773455888032913, + 0.027666568756103516, + 0.045300740748643875, + 0.03623325005173683, + -0.046061139553785324, + -0.00023302776389755309, + 0.010822237469255924, + 0.005396110936999321, + -0.017266767099499702, + -0.042832016944885254, + 0.000606586632784456, + 0.0053397659212350845, + 0.0031394879333674908, + -0.014797745272517204, + 0.04381722956895828, + 0.014831854961812496, + -0.028363173827528954, + -0.03383432328701019, + 0.031214935705065727, + 0.08290299773216248, + 0.02584749460220337, + 0.03909004107117653, + 0.0349324494600296, + -0.02809840254485607, + -0.05105957016348839, + -0.003338165581226349, + -0.044679734855890274, + -0.016897771507501602, + 0.04674547165632248, + -0.007903410121798515, + -0.06998921185731888, + -0.00042846743599511683, + 0.024417705833911896, + -0.025209281593561172, + 0.01556653156876564, + 0.014268666505813599, + 0.003608678700402379, + -0.03939566761255264, + -0.06773681193590164, + -0.006201360374689102, + -0.038590289652347565, + -0.04736294224858284, + 0.054148513823747635, + -0.04606981575489044, + -0.0216580331325531, + -0.02679860219359398, + -0.0387553870677948, + -0.06344524025917053, + 0.03479386121034622, + -0.03733435273170471, + -0.02730160392820835, + 0.033917054533958435, + -0.01180818397551775, + 0.017153600230813026, + -0.017416737973690033, + -0.05760662630200386, + 0.005898391827940941, + -0.06530486792325974, + -0.044563937932252884, + 0.013888120651245117, + -0.05623237416148186, + 0.012411084957420826, + 0.031150581315159798, + -0.05592577904462814, + 0.00470947427675128, + -0.012832939624786377, + -0.051747485995292664, + 0.08193240314722061, + 0.04231006279587746, + 0.060262277722358704, + -0.05159112438559532, + -0.0050200470723211765, + -0.028430450707674026, + 0.048510000109672546, + -0.03564736619591713, + 0.040754739195108414, + 0.028905360028147697, + 0.038992658257484436, + -0.0016340425936505198, + -0.008609502576291561, + 0.006150763016194105, + -0.004917451646178961, + -0.073238305747509, + 0.014989398419857025, + 0.039281148463487625, + 0.026978103443980217, + -0.01927020587027073, + 0.03466196730732918, + -0.026035647839307785, + 0.04892418906092644, + -0.021652784198522568, + -0.05802053213119507, + -0.007418050896376371, + 0.0004157012444920838, + -0.003113806713372469, + -0.04071711003780365, + -0.03618207946419716, + 0.02473904751241207, + 0.005009343847632408, + -0.003717175917699933, + 0.06939894706010818, + -0.03745422139763832, + -0.056340575218200684, + 0.020305491983890533, + 0.01624903455376625, + -0.0414578877389431, + -0.006078140810132027, + -0.00001147667171608191, + -0.03432951495051384, + 0.04420928284525871, + 0.03278480842709541, + 0.0019583101384341717, + 0.006309302523732185, + -0.041066475212574005, + 0.0045809498988091946, + 0.03785176947712898, + -0.030005238950252533, + 0.0525495707988739, + -0.024305732920765877, + 0.0009916236158460379, + 0.026999732479453087, + -0.013087441213428974, + 0.0179459061473608, + 0.029477974399924278, + -0.0367741696536541, + 0.06801317632198334, + 0.0028329314664006233, + 0.02241651341319084, + -0.012403511442244053, + -0.06610151380300522, + -0.01063673384487629, + -0.0017124409787356853, + 0.009534181095659733, + 0.040957942605018616, + 0.030794164165854454, + -0.03456086665391922, + -0.04005061089992523, + -0.053644463419914246, + 0.011436519213020802, + -0.011531087569892406, + -0.020456653088331223, + -0.0051927464082837105, + 0.008158964104950428, + 0.02680368535220623, + 0.0762314721941948, + -0.022566702216863632, + 0.03298712149262428, + 0.0227197352796793, + 0.01228597667068243, + 0.0397644080221653, + -0.035839907824993134, + 0.04033581539988518, + -0.023282667621970177, + 0.0071014040149748325, + 0.060574986040592194, + 0.03509609028697014, + -0.023195691406726837, + 0.02417423389852047 + ], + "index": 0, + "object": "embedding" + } + ], + "model": "text-embedding-004", + "object": "list", + "usage": { + "completion_tokens": 0, + "prompt_tokens": 7, + "total_tokens": 7 + } + } + } + } +] \ No newline at end of file diff --git a/tests/vertexai_integration_test.rs b/tests/vertexai_integration_test.rs new file mode 100644 index 0000000..daddc0b --- /dev/null +++ b/tests/vertexai_integration_test.rs @@ -0,0 +1,415 @@ +use hub::models::chat::{ChatCompletionRequest, ChatCompletionResponse}; +use hub::models::content::{ChatCompletionMessage, ChatMessageContent}; +use hub::models::embeddings::{EmbeddingsInput, EmbeddingsRequest}; +use hub::models::tool_definition::{FunctionDefinition, ToolDefinition}; +use hub::providers::provider::Provider; +use hub::providers::vertexai::VertexAIProvider; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::env; +use std::fs; +use std::path::PathBuf; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const RECORDINGS_DIR: &str = "tests/recordings/vertexai"; + +#[derive(serde::Serialize, serde::Deserialize)] +struct RecordedInteraction { + request: RequestData, + response: ResponseData, +} + +#[derive(serde::Serialize, serde::Deserialize)] +struct RequestData { + method: wiremock::http::Method, + path: String, + headers: HashMap, + body: Option, +} + +#[derive(serde::Serialize, serde::Deserialize)] +struct ResponseData { + status: u16, + body: Value, +} + +struct TestConfig { + project_id: String, + location: String, +} + +impl TestConfig { + fn new() -> Self { + if env::var("RECORD").is_ok() { + dotenv::from_filename(".env").ok(); + Self { + project_id: env::var("VERTEX_PROJECT_ID") + .expect("VERTEX_PROJECT_ID must be set for recording"), + location: env::var("VERTEX_LOCATION").unwrap_or_else(|_| "us-central1".to_string()), + } + } else { + Self::from_recordings().unwrap_or_else(|| Self { + project_id: "extended-legend-445620-u1".to_string(), + location: "us-central1".to_string(), + }) + } + } + + fn from_recordings() -> Option { + let recordings_dir = PathBuf::from(RECORDINGS_DIR); + if let Ok(content) = fs::read_to_string(recordings_dir.join("chat_completion.json")) { + if let Ok(interactions) = serde_json::from_str::>(&content) { + if let Some(interaction) = interactions.first() { + if let Some((project_id, location)) = + Self::extract_from_path(&interaction.request.path) + { + return Some(Self { + project_id: project_id.to_string(), + location: location.to_string(), + }); + } + } + } + } + None + } + + fn extract_from_path(path: &str) -> Option<(&str, &str)> { + let parts: Vec<&str> = path.split('/').collect(); + if parts.len() >= 7 { + Some((parts[3], parts[5])) + } else { + None + } + } +} + +async fn setup_mock_server(test_name: &str) -> MockServer { + let mock_server = MockServer::start().await; + let recordings_dir = PathBuf::from(RECORDINGS_DIR); + + if env::var("RECORD").is_err() { + if let Ok(content) = fs::read_to_string(recordings_dir.join(format!("{}.json", test_name))) + { + if let Ok(interactions) = serde_json::from_str::>(&content) { + for interaction in interactions { + setup_mock(&mock_server, interaction).await; + } + } + } + } + + mock_server +} + +async fn setup_mock(mock_server: &MockServer, interaction: RecordedInteraction) { + let mut mock = + Mock::given(method(interaction.request.method)).and(path(&interaction.request.path)); + + for (key, value) in interaction.request.headers { + mock = mock.and(wiremock::matchers::header(key.as_str(), value.as_str())); + } + + mock.respond_with( + ResponseTemplate::new(interaction.response.status).set_body_json(interaction.response.body), + ) + .mount(mock_server) + .await; +} + +async fn save_interaction(test_name: &str, interaction: RecordedInteraction) { + let recordings_dir = PathBuf::from(RECORDINGS_DIR); + fs::create_dir_all(&recordings_dir).unwrap_or_default(); + let recording_path = recordings_dir.join(format!("{}.json", test_name)); + + let mut interactions = Vec::new(); + if let Ok(content) = fs::read_to_string(&recording_path) { + if let Ok(mut existing) = serde_json::from_str::>(&content) { + interactions.append(&mut existing); + } + } + + interactions.push(interaction); + + if let Ok(content) = serde_json::to_string_pretty(&interactions) { + fs::write(&recording_path, content).expect("Failed to save recording"); + } +} + +async fn create_test_provider(mock_server: &MockServer) -> VertexAIProvider { + let config = TestConfig::new(); + let mut params = HashMap::new(); + params.insert("project_id".to_string(), config.project_id); + params.insert("location".to_string(), config.location); + + if env::var("RECORD").is_err() { + params.insert("test_base_url".to_string(), mock_server.uri()); + params.insert("skip_authentication".to_string(), "true".to_string()); + } + + let provider_config = hub::config::models::Provider { + key: "vertexai".to_string(), + r#type: "vertexai".to_string(), + api_key: String::new(), + params, + }; + + VertexAIProvider::new(&provider_config) +} + +async fn record_chat_response(response: &ChatCompletionResponse) -> Value { + match response { + ChatCompletionResponse::NonStream(completion) => { + serde_json::to_value(completion).unwrap_or_else(|_| json!(null)) + } + ChatCompletionResponse::Stream(_) => { + json!({ + "type": "stream", + "status": "recorded" + }) + } + } +} + +#[tokio::test] +async fn test_chat_completion() { + let mock_server = setup_mock_server("chat_completion").await; + let provider = create_test_provider(&mock_server).await; + let config = TestConfig::new(); + + let model_config = hub::config::models::ModelConfig { + key: "gemini-pro".to_string(), + r#type: "gemini-pro".to_string(), + provider: "vertexai".to_string(), + params: HashMap::new(), + }; + + let request = ChatCompletionRequest { + model: "gemini-pro".to_string(), + messages: vec![ChatCompletionMessage { + role: "user".to_string(), + content: Some(ChatMessageContent::String( + "What is the capital of France?".to_string(), + )), + name: None, + tool_calls: None, + }], + temperature: Some(0.7), + stream: None, + max_tokens: Some(100), + top_p: None, + n: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + tool_choice: None, + tools: None, + user: None, + parallel_tool_calls: None, + }; + + let response = provider + .chat_completions(request.clone(), &model_config) + .await; + + if env::var("RECORD").is_ok() { + if let Ok(resp) = &response { + let response_value = match resp { + ChatCompletionResponse::NonStream(completion) => { + serde_json::to_value(completion).unwrap() + } + ChatCompletionResponse::Stream(_) => { + json!({ + "type": "stream", + "status": "recorded" + }) + } + }; + + let interaction = RecordedInteraction { + request: RequestData { + method: wiremock::http::Method::Post, + path: format!( + "/v1/projects/{}/locations/{}/publishers/google/models/gemini-pro:generateContent", + config.project_id, config.location + ), + headers: HashMap::new(), + body: Some(serde_json::to_value(&request).unwrap()), + }, + response: ResponseData { + status: 200, + body: response_value, + }, + }; + save_interaction("chat_completion", interaction).await; + } + } + + assert!(response.is_ok(), "Chat completion request failed"); + if let Ok(ChatCompletionResponse::NonStream(completion)) = response { + assert!(!completion.choices.is_empty(), "No choices in response"); + assert!( + completion.choices[0].message.content.is_some(), + "No content in response" + ); + } +} + +#[tokio::test] +async fn test_chat_completion_with_tools() { + let mock_server = setup_mock_server("chat_completion_with_tools").await; + let provider = create_test_provider(&mock_server).await; + + let model_config = hub::config::models::ModelConfig { + key: "gemini-pro".to_string(), + r#type: "gemini-pro".to_string(), + provider: "vertexai".to_string(), + params: HashMap::new(), + }; + + let request = ChatCompletionRequest { + model: "gemini-pro".to_string(), + messages: vec![ChatCompletionMessage { + role: "user".to_string(), + content: Some(ChatMessageContent::String( + "What's the weather in San Francisco?".to_string(), + )), + name: None, + tool_calls: None, + }], + temperature: Some(0.7), + stream: None, + max_tokens: Some(100), + top_p: None, + n: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + tool_choice: None, + tools: Some(vec![ToolDefinition { + tool_type: "function".to_string(), + function: FunctionDefinition { + name: "get_weather".to_string(), + description: Some("Get the current weather in a location".to_string()), + parameters: Some( + serde_json::from_value(json!({ + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get weather for" + } + }, + "required": ["location"] + })) + .unwrap(), + ), + strict: None, + }, + }]), + user: None, + parallel_tool_calls: None, + }; + + let response = provider + .chat_completions(request.clone(), &model_config) + .await; + + if env::var("RECORD").is_ok() { + if let Ok(resp) = &response { + let interaction = RecordedInteraction { + request: RequestData { + method: wiremock::http::Method::Post, + path: "gemini-pro:generateContent".to_string(), + headers: HashMap::new(), + body: Some(serde_json::to_value(&request).unwrap()), + }, + response: ResponseData { + status: 200, + body: record_chat_response(resp).await, + }, + }; + save_interaction("chat_completion_with_tools", interaction).await; + } + } + + assert!( + response.is_ok(), + "Chat completion with tools request failed" + ); + + if let Ok(ChatCompletionResponse::NonStream(completion)) = response { + assert!(!completion.choices.is_empty(), "No choices in response"); + let tool_calls = completion.choices[0].message.tool_calls.as_ref(); + assert!(tool_calls.is_some(), "No tool calls in response"); + + if let Some(tool_calls) = tool_calls { + assert!(!tool_calls.is_empty(), "Empty tool calls"); + assert_eq!( + tool_calls[0].function.name, "get_weather", + "Incorrect function name" + ); + + let args: Value = serde_json::from_str(&tool_calls[0].function.arguments).unwrap(); + assert!(args["location"].is_string(), "Location should be a string"); + } + } +} + +#[tokio::test] +async fn test_embeddings_functionality() { + println!("Starting embeddings test"); + let mock_server = setup_mock_server("embeddings_test").await; + let provider = create_test_provider(&mock_server).await; + + let model_config = hub::config::models::ModelConfig { + key: "textembedding-gecko".to_string(), + r#type: "textembedding-gecko".to_string(), + provider: "vertexai".to_string(), + params: HashMap::new(), + }; + + let request = EmbeddingsRequest { + model: "textembedding-gecko".to_string(), + input: EmbeddingsInput::Single("Explore embeddings test.".to_string()), + user: None, + encoding_format: None, + }; + + let response = provider.embeddings(request.clone(), &model_config).await; + + if env::var("RECORD").is_ok() { + if let Ok(resp) = &response { + let config = TestConfig::new(); + let interaction = RecordedInteraction { + request: RequestData { + method: wiremock::http::Method::Post, + path: format!( + "/v1/projects/{}/locations/{}/publishers/google/models/textembedding-gecko:predict", + config.project_id, config.location + ), + headers: HashMap::new(), + body: Some(serde_json::to_value(&request).unwrap()), + }, + response: ResponseData { + status: 200, + body: serde_json::to_value(resp).unwrap(), + }, + }; + save_interaction("embeddings_test", interaction).await; + } + } + + assert!(response.is_ok(), "Embeddings request failed"); + + if let Ok(embeddings) = response { + assert!(!embeddings.data.is_empty(), "Embeddings response is empty"); + assert!( + !embeddings.data[0].embedding.is_empty(), + "Embedding vector is empty" + ); + } +}