use native client for inference in UI (tensorzero#2975)

* added ClientInferenceParams and inferenceResponse to TS bindings

* refactored getting native tensorzero client into its own file

* refactored utilites for inference

* fixed import issues

* fixed bindings generation

* implemented a rust-safe stringify helper

* implemented type safe conversion to rust input type

* more typing fixes

* inference works via rust client for server inferences in playground

* deprecated ts inference client

* cleaned up PR

* consolidated building bindings into cargo tsbuild

* added a comment about caching settings in playground

* inference should throw if `extra_body` is set

* fixed check for extra body

* fixed error handling for try with variant

* fixed issue with merge

Author: virajmehta

.cargo/config.toml

@@ -44,5 +44,11 @@ build-e2e = "build --bin gateway --features e2e_tests"
 run-e2e = "run --bin gateway --features e2e_tests -- --config-file tensorzero-core/tests/e2e/tensorzero.toml"
 watch-e2e = "watch -x run-e2e"
 
-# Export Typescript bindings for TensorZero
-tsbuild = ["test", "export_bindings", "-p", "tensorzero-core"]
+tsbuild = [
+    "test",
+    "export_bindings",
+    "-p",
+    "tensorzero-core",
+    "-p",
+    "tensorzero",
+] # Export Typescript bindings for TensorZero

Cargo.lock

@@ -21,6 +21,7 @@ reqwest-eventsource = { workspace = true }
 async-stream = { workspace = true }
 tokio-stream = { workspace = true }
 tensorzero-core = { path = "../../tensorzero-core" }
+ts-rs = { workspace = true }
 url = { workspace = true }
 thiserror = "2.0.11"
 pyo3 = { workspace = true, optional = true }

clients/rust/Cargo.toml

@@ -21,7 +21,8 @@ use crate::client_input::{test_client_input_to_input, ClientInput};
 // This is a copy-paste of the `Params` struct from `tensorzero_core::endpoints::inference::Params`.
 // with just the `credentials` field adjusted to allow serialization.
 /// The expected payload is a JSON object with the following fields:
-#[derive(Clone, Debug, Serialize, Default)]
+#[derive(Clone, Debug, Deserialize, Serialize, Default, ts_rs::TS)]
+#[ts(export)]
 pub struct ClientInferenceParams {
     // The function name. Exactly one of `function_name` or `model_name` must be provided.
     pub function_name: Option<String>,
@@ -61,16 +62,27 @@ pub struct ClientInferenceParams {
     // If provided for a JSON inference, the inference will use the specified output schema instead of the
     // configured one. We only lazily validate this schema.
     pub output_schema: Option<Value>,
+    #[ts(type = "Map<string, string>")]
     pub credentials: HashMap<String, ClientSecretString>,
     pub cache_options: CacheParamsOptions,
     /// If `true`, add an `original_response` field to the response, containing the raw string response from the model.
     /// Note that for complex variants (e.g. `experimental_best_of_n_sampling`), the response may not contain `original_response`
     /// if the fuser/judge model failed
     #[serde(default)]
     pub include_original_response: bool,
+    // NOTE: Currently, ts_rs does not handle #[serde(transparent)] correctly,
+    // so we disable the type generation for the extra_body and extra_headers fields.
+    // I tried doing a direct #[ts(type = "InferenceExtraBody[]")] and
+    // a #[ts(as = "Vec<InferenceExtraBody>")] and these would generate the types but then
+    // type checking would fail because the ClientInferenceParams struct would not be
+    // generated with the correct import.
+    //
+    // Not sure if this is solvable with the existing crate.
     #[serde(default)]
+    #[ts(skip)]
     pub extra_body: UnfilteredInferenceExtraBody,
     #[serde(default)]
+    #[ts(skip)]
     pub extra_headers: UnfilteredInferenceExtraHeaders,
     pub internal_dynamic_variant_config: Option<UninitializedVariantInfo>,
 }

clients/rust/src/client_inference_params.rs

@@ -11,6 +11,8 @@ use tensorzero_derive::TensorZeroDeserialize;
 // Like the normal `Input` type, but with `ClientInputMessage` instead of `InputMessage`.
 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Default)]
 #[serde(deny_unknown_fields)]
+#[derive(ts_rs::TS)]
+#[ts(export)]
 pub struct ClientInput {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub system: Option<Value>,
@@ -21,6 +23,8 @@ pub struct ClientInput {
 // Like the normal `InputMessage` type, but with `ClientInputMessageContent` instead of `InputMessageContent`.
 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 #[serde(deny_unknown_fields)]
+#[derive(ts_rs::TS)]
+#[ts(export)]
 pub struct ClientInputMessage {
     pub role: Role,
     #[serde(deserialize_with = "deserialize_content")]
@@ -30,6 +34,8 @@ pub struct ClientInputMessage {
 #[derive(Clone, Debug, TensorZeroDeserialize, Serialize, PartialEq)]
 #[serde(tag = "type")]
 #[serde(rename_all = "snake_case")]
+#[derive(ts_rs::TS)]
+#[ts(export)]
 pub enum ClientInputMessageContent {
     Text(TextKind),
     ToolCall(ToolCallInput),

clients/rust/src/client_input.rs

@@ -13,7 +13,7 @@ mkdir -p lib/bindings
 # Generate TypeScript bindings from Rust code
 echo "Generating TypeScript bindings from Rust..."
 cd ../..
-TS_RS_EXPORT_DIR="../internal/tensorzero-node/lib/bindings" cargo tsbuild
+TS_RS_EXPORT_DIR="$(pwd)/internal/tensorzero-node/lib/bindings" cargo tsbuild
 cd internal/tensorzero-node
 
 # Generate index file

internal/tensorzero-node/build-bindings.sh

@@ -1,6 +1,4 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
-import type { ExtraBodyConfig } from "./ExtraBodyConfig";
-import type { ExtraHeadersConfig } from "./ExtraHeadersConfig";
 import type { JsonMode } from "./JsonMode";
 import type { PathWithContents } from "./PathWithContents";
 import type { RetryConfig } from "./RetryConfig";
@@ -20,6 +18,4 @@ export type BestOfNEvaluatorConfig = {
   stop_sequences: Array<string> | null;
   json_mode: JsonMode | null;
   retries: RetryConfig;
-  extra_body: ExtraBodyConfig | null;
-  extra_headers: ExtraHeadersConfig | null;
 };

internal/tensorzero-node/lib/bindings/BestOfNEvaluatorConfig.ts

@@ -1,3 +1,3 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
 
-export type ExtraHeader = { name: string } & ({ value: string } | "delete");
+export type CacheEnabledMode = "on" | "off" | "read_only" | "write_only";

…sorzero-node/lib/bindings/ExtraHeader.ts …ro-node/lib/bindings/CacheEnabledMode.tsinternal/tensorzero-node/lib/bindings/ExtraHeader.ts renamed to internal/tensorzero-node/lib/bindings/CacheEnabledMode.ts internal/tensorzero-node/lib/bindings/ExtraHeader.ts renamed to internal/tensorzero-node/lib/bindings/CacheEnabledMode.ts

@@ -0,0 +1,7 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { CacheEnabledMode } from "./CacheEnabledMode";
+
+export type CacheParamsOptions = {
+  max_age_s: number | null;
+  enabled: CacheEnabledMode;
+};

internal/tensorzero-node/lib/bindings/CacheParamsOptions.ts

@@ -1,6 +1,4 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
-import type { ExtraBodyConfig } from "./ExtraBodyConfig";
-import type { ExtraHeadersConfig } from "./ExtraHeadersConfig";
 import type { JsonMode } from "./JsonMode";
 import type { PathWithContents } from "./PathWithContents";
 import type { RetryConfig } from "./RetryConfig";
@@ -20,6 +18,4 @@ export type ChainOfThoughtConfig = {
   stop_sequences: Array<string> | null;
   json_mode: JsonMode | null;
   retries: RetryConfig;
-  extra_body: ExtraBodyConfig | null;
-  extra_headers: ExtraHeadersConfig | null;
 };