Skip to content

Commit b16b88a

Browse files
hongantran3804hongantran3804
committed
Merge branch 'main' of https://github.com/pycoder49/tensorzero into HongAn
2 parents bd4a309 + ff0e58a commit b16b88a

File tree

62 files changed

+1305
-741
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

62 files changed

+1305
-741
lines changed

‎.github/workflows/general.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,9 @@ jobs:
481481
- name: Download ClickHouse fixtures
482482
run: uv run ./ui/fixtures/download-fixtures.py
483483

484+
- name: Cleanup disk space
485+
run: ./ci/free-disk-space.sh
486+
484487
- name: Set up TENSORZERO_CLICKHOUSE_URL for E2E tests
485488
run: |
486489
echo "TENSORZERO_CLICKHOUSE_URL=http://chuser:chpassword@localhost:8123/tensorzero_e2e_tests" >> $GITHUB_ENV

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,11 @@ print_stderr = "deny"
8383
print_stdout = "deny"
8484
redundant_closure_for_method_calls = "deny"
8585
todo = "deny"
86+
trivially_copy_pass_by_ref = "deny"
8687
unimplemented = "deny"
8788
uninlined_format_args = "deny"
8889
unreachable = "deny"
90+
unused_self = "deny"
8991
unwrap_used = "deny"
9092

9193
[profile.performance]

‎clients/python/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -576,6 +576,7 @@ impl TensorZeroGateway {
576576
}
577577

578578
/// Close the connection to the TensorZero gateway.
579+
#[expect(clippy::unused_self)]
579580
fn close(&self) {
580581
// TODO - implement closing the 'reqwest' connection pool: https://github.com/tensorzero/tensorzero/issues/857
581582
}

‎clients/rust/src/lib.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1314,13 +1314,10 @@ mod tests {
13141314
use tracing_test::traced_test;
13151315

13161316
#[tokio::test]
1317-
#[ignore] // TODO - set an environment variable, or create a new config with dummy credentials
13181317
async fn test_missing_clickhouse() {
13191318
// This config file requires ClickHouse, so it should fail if no ClickHouse URL is provided
13201319
let err = ClientBuilder::new(ClientBuilderMode::EmbeddedGateway {
1321-
config_file: Some(PathBuf::from(
1322-
"../../examples/haiku-hidden-preferences/config/tensorzero.toml",
1323-
)),
1320+
config_file: Some(PathBuf::from("tests/test_config.toml")),
13241321
clickhouse_url: None,
13251322
timeout: None,
13261323
verify_credentials: true,
@@ -1330,7 +1327,7 @@ mod tests {
13301327
.expect_err("ClientBuilder should have failed");
13311328
let err_msg = err.to_string();
13321329
assert!(
1333-
err.to_string().contains("Missing ClickHouse URL"),
1330+
err_msg.contains("Missing environment variable TENSORZERO_CLICKHOUSE_URL"),
13341331
"Bad error message: {err_msg}"
13351332
);
13361333
}

‎clients/rust/tests/templates/system_template.minijinja‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
You are a helpful assistant.

‎clients/rust/tests/templates/user_template.minijinja‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{{ user_text }}

‎clients/rust/tests/test_config.toml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Test configuration for testing ClickHouse requirement
2+
# This config explicitly enables observability, which requires ClickHouse
3+
4+
[gateway.observability]
5+
enabled = true
6+
7+
# Minimal function definition for a valid config
8+
[functions.test_function]
9+
type = "chat"
10+
11+
[functions.test_function.variants.test_variant]
12+
type = "chat_completion"
13+
model = "openai::gpt-4o-mini-2024-07-18"
14+
system_template = "templates/system_template.minijinja"
15+
user_template = "templates/user_template.minijinja"

‎internal/tensorzero-node/lib/bindings/TomlRelativePath.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,4 @@
66
* all paths (e.g. `system_schema`) as `TomlRelativePath`s, which will
77
* track the original `.toml` file in order to perform correct relative path resolution.
88
*/
9-
export type TomlRelativePath = { path: string };
9+
export type TomlRelativePath = { __tensorzero_remapped_path: string };

‎tensorzero-core/src/config_parser/mod.rs‎

Lines changed: 52 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use crate::variant::dicl::UninitializedDiclConfig;
3939
use crate::variant::mixture_of_n::UninitializedMixtureOfNConfig;
4040
use crate::variant::{Variant, VariantConfig, VariantInfo};
4141
use std::error::Error as StdError;
42+
use toml::de::DeTable;
4243

4344
pub mod gateway;
4445
pub mod path;
@@ -414,15 +415,6 @@ impl Config {
414415
config_path: &Path,
415416
validate_credentials: bool,
416417
) -> Result<Config, Error> {
417-
let config_table = match UninitializedConfig::read_toml_config(config_path)? {
418-
Some(table) => table,
419-
None => {
420-
return Err(ErrorDetails::Config {
421-
message: format!("Config file not found: {config_path:?}"),
422-
}
423-
.into())
424-
}
425-
};
426418
let base_path = match PathBuf::from(&config_path).parent() {
427419
Some(base_path) => base_path.to_path_buf(),
428420
None => {
@@ -434,6 +426,15 @@ impl Config {
434426
.into());
435427
}
436428
};
429+
let config_table = match UninitializedConfig::read_toml_config(config_path, &base_path)? {
430+
Some(table) => table,
431+
None => {
432+
return Err(ErrorDetails::Config {
433+
message: format!("Config file not found: {config_path:?}"),
434+
}
435+
.into())
436+
}
437+
};
437438
let config = if cfg!(feature = "e2e_tests") || !validate_credentials {
438439
SKIP_CREDENTIAL_VALIDATION
439440
.scope((), Self::load_from_toml(config_table, base_path))
@@ -462,17 +463,13 @@ impl Config {
462463
let functions = uninitialized_config
463464
.functions
464465
.into_iter()
465-
.map(|(name, config)| config.load(&name, &base_path).map(|c| (name, Arc::new(c))))
466+
.map(|(name, config)| config.load(&name).map(|c| (name, Arc::new(c))))
466467
.collect::<Result<HashMap<String, Arc<FunctionConfig>>, Error>>()?;
467468

468469
let tools = uninitialized_config
469470
.tools
470471
.into_iter()
471-
.map(|(name, config)| {
472-
config
473-
.load(&base_path, name.clone())
474-
.map(|c| (name, Arc::new(c)))
475-
})
472+
.map(|(name, config)| config.load(name.clone()).map(|c| (name, Arc::new(c))))
476473
.collect::<Result<HashMap<String, Arc<StaticToolConfig>>, Error>>()?;
477474

478475
let models = uninitialized_config
@@ -552,7 +549,7 @@ impl Config {
552549
let mut evaluations = HashMap::new();
553550
for (name, evaluation_config) in uninitialized_config.evaluations {
554551
let (evaluation_config, evaluation_function_configs, evaluation_metric_configs) =
555-
evaluation_config.load(&config.functions, &base_path, &name)?;
552+
evaluation_config.load(&config.functions, &name)?;
556553
evaluations.insert(name, Arc::new(EvaluationConfig::Static(evaluation_config)));
557554
for (evaluation_function_name, evaluation_function_config) in
558555
evaluation_function_configs
@@ -832,9 +829,9 @@ impl FunctionsConfigPyClass {
832829
}
833830
}
834831

835-
/// A trait for loading configs with a base path
832+
/// A trait for loading configs
836833
pub trait LoadableConfig<T> {
837-
fn load<P: AsRef<Path>>(self, base_path: P) -> Result<T, Error>;
834+
fn load(self) -> Result<T, Error>;
838835
}
839836

840837
/// This struct is used to deserialize the TOML config file
@@ -879,28 +876,26 @@ pub struct ProviderTypesConfig {
879876

880877
impl UninitializedConfig {
881878
/// Read a file from the file system and parse it as TOML
882-
fn read_toml_config(path: &Path) -> Result<Option<toml::Table>, Error> {
879+
fn read_toml_config(path: &Path, base_path: &Path) -> Result<Option<toml::Table>, Error> {
883880
if !path.exists() {
884881
return Ok(None);
885882
}
886-
Ok(Some(
887-
std::fs::read_to_string(path)
888-
.map_err(|_| {
889-
Error::new(ErrorDetails::Config {
890-
message: format!("Failed to read config file: {}", path.to_string_lossy()),
891-
})
892-
})?
893-
.parse::<toml::Table>()
894-
.map_err(|e| {
895-
Error::new(ErrorDetails::Config {
896-
message: format!(
897-
"Failed to parse config file `{}` as valid TOML: {}",
898-
path.to_string_lossy(),
899-
e
900-
),
901-
})
902-
})?,
903-
))
883+
let contents = std::fs::read_to_string(path).map_err(|_| {
884+
Error::new(ErrorDetails::Config {
885+
message: format!("Failed to read config file: {}", path.to_string_lossy()),
886+
})
887+
})?;
888+
let table = DeTable::parse(&contents).map_err(|e| {
889+
Error::new(ErrorDetails::Config {
890+
message: format!(
891+
"Failed to parse config file `{}` as valid TOML: {}",
892+
path.to_string_lossy(),
893+
e
894+
),
895+
})
896+
})?;
897+
let table = path::resolve_toml_relative_paths(table.into_inner(), base_path)?;
898+
Ok(Some(table))
904899
}
905900
}
906901

@@ -960,29 +955,25 @@ struct UninitializedFunctionConfigJson {
960955
}
961956

962957
impl UninitializedFunctionConfig {
963-
pub fn load<P: AsRef<Path>>(
964-
self,
965-
function_name: &str,
966-
base_path: P,
967-
) -> Result<FunctionConfig, Error> {
958+
pub fn load(self, function_name: &str) -> Result<FunctionConfig, Error> {
968959
match self {
969960
UninitializedFunctionConfig::Chat(params) => {
970961
let system_schema = params
971962
.system_schema
972-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
963+
.map(StaticJSONSchema::from_path)
973964
.transpose()?;
974965
let user_schema = params
975966
.user_schema
976-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
967+
.map(StaticJSONSchema::from_path)
977968
.transpose()?;
978969
let assistant_schema = params
979970
.assistant_schema
980-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
971+
.map(StaticJSONSchema::from_path)
981972
.transpose()?;
982973
let variants = params
983974
.variants
984975
.into_iter()
985-
.map(|(name, variant)| variant.load(&base_path).map(|v| (name, Arc::new(v))))
976+
.map(|(name, variant)| variant.load().map(|v| (name, Arc::new(v))))
986977
.collect::<Result<HashMap<_, _>, Error>>()?;
987978
for (name, variant) in variants.iter() {
988979
if let VariantConfig::ChatCompletion(chat_config) = &variant.inner {
@@ -1010,26 +1001,26 @@ impl UninitializedFunctionConfig {
10101001
UninitializedFunctionConfig::Json(params) => {
10111002
let system_schema = params
10121003
.system_schema
1013-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
1004+
.map(StaticJSONSchema::from_path)
10141005
.transpose()?;
10151006
let user_schema = params
10161007
.user_schema
1017-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
1008+
.map(StaticJSONSchema::from_path)
10181009
.transpose()?;
10191010
let assistant_schema = params
10201011
.assistant_schema
1021-
.map(|path| StaticJSONSchema::from_path(path, base_path.as_ref()))
1012+
.map(StaticJSONSchema::from_path)
10221013
.transpose()?;
10231014
let output_schema = match params.output_schema {
1024-
Some(path) => StaticJSONSchema::from_path(path, base_path.as_ref())?,
1015+
Some(path) => StaticJSONSchema::from_path(path)?,
10251016
None => StaticJSONSchema::default(),
10261017
};
10271018
let implicit_tool_call_config =
10281019
create_implicit_tool_call_config(output_schema.clone());
10291020
let variants = params
10301021
.variants
10311022
.into_iter()
1032-
.map(|(name, variant)| variant.load(&base_path).map(|v| (name, Arc::new(v))))
1023+
.map(|(name, variant)| variant.load().map(|v| (name, Arc::new(v))))
10331024
.collect::<Result<HashMap<_, _>, Error>>()?;
10341025

10351026
for (name, variant) in variants.iter() {
@@ -1107,22 +1098,20 @@ pub enum UninitializedVariantConfig {
11071098
}
11081099

11091100
impl UninitializedVariantInfo {
1110-
pub fn load<P: AsRef<Path>>(self, base_path: P) -> Result<VariantInfo, Error> {
1101+
pub fn load(self) -> Result<VariantInfo, Error> {
11111102
let inner = match self.inner {
11121103
UninitializedVariantConfig::ChatCompletion(params) => {
1113-
VariantConfig::ChatCompletion(params.load(base_path)?)
1104+
VariantConfig::ChatCompletion(params.load()?)
11141105
}
11151106
UninitializedVariantConfig::BestOfNSampling(params) => {
1116-
VariantConfig::BestOfNSampling(params.load(base_path)?)
1117-
}
1118-
UninitializedVariantConfig::Dicl(params) => {
1119-
VariantConfig::Dicl(params.load(base_path)?)
1107+
VariantConfig::BestOfNSampling(params.load()?)
11201108
}
1109+
UninitializedVariantConfig::Dicl(params) => VariantConfig::Dicl(params.load()?),
11211110
UninitializedVariantConfig::MixtureOfN(params) => {
1122-
VariantConfig::MixtureOfN(params.load(base_path)?)
1111+
VariantConfig::MixtureOfN(params.load()?)
11231112
}
11241113
UninitializedVariantConfig::ChainOfThought(params) => {
1125-
VariantConfig::ChainOfThought(params.load(base_path)?)
1114+
VariantConfig::ChainOfThought(params.load()?)
11261115
}
11271116
};
11281117
Ok(VariantInfo {
@@ -1143,12 +1132,8 @@ pub struct UninitializedToolConfig {
11431132
}
11441133

11451134
impl UninitializedToolConfig {
1146-
pub fn load<P: AsRef<Path>>(
1147-
self,
1148-
base_path: P,
1149-
name: String,
1150-
) -> Result<StaticToolConfig, Error> {
1151-
let parameters = StaticJSONSchema::from_path(self.parameters, base_path.as_ref())?;
1135+
pub fn load(self, name: String) -> Result<StaticToolConfig, Error> {
1136+
let parameters = StaticJSONSchema::from_path(self.parameters)?;
11521137
Ok(StaticToolConfig {
11531138
name: self.name.unwrap_or(name),
11541139
description: self.description,
@@ -1168,15 +1153,8 @@ pub struct PathWithContents {
11681153
}
11691154

11701155
impl PathWithContents {
1171-
pub fn from_path<P: AsRef<Path>>(
1172-
path: TomlRelativePath,
1173-
base_path: Option<P>,
1174-
) -> Result<Self, Error> {
1175-
let full_path = if let Some(base_path) = base_path.as_ref() {
1176-
&base_path.as_ref().join(path.path())
1177-
} else {
1178-
path.path()
1179-
};
1156+
pub fn from_path(path: TomlRelativePath) -> Result<Self, Error> {
1157+
let full_path = path.path();
11801158
let contents = std::fs::read_to_string(full_path).map_err(|e| {
11811159
Error::new(ErrorDetails::Config {
11821160
message: format!(

0 commit comments

Comments
 (0)