feat(test-util): construct field-based arg schemas

Author: lemorage

src/ops/functions/embed_text.rs

@@ -104,7 +104,6 @@ mod tests {
     #[tokio::test]
     #[ignore = "This test requires OpenAI API key or a configured local LLM and may make network calls."]
     async fn test_embed_text() {
-        // Using OpenAI as an example.
         let spec = Spec {
             api_type: LlmApiType::OpenAi,
             model: "text-embedding-ada-002".to_string(),
@@ -118,11 +117,7 @@ mod tests {
 
         let input_args_values = vec![text_content.to_string().into()];
 
-        let input_arg_schemas = vec![build_arg_schema(
-            "text",
-            text_content.to_string().into(),
-            BasicValueType::Str,
-        )];
+        let input_arg_schemas = vec![build_arg_schema("text", BasicValueType::Str)];
 
         let result = test_flow_function(factory, spec, input_arg_schemas, input_args_values).await;
 

src/ops/functions/extract_by_llm.rs

@@ -185,7 +185,6 @@ mod tests {
             attrs: Arc::new(BTreeMap::new()),
         };
 
-        // Spec using OpenAI as an example.
         let spec = Spec {
             llm_spec: LlmSpec {
                 api_type: crate::llm::LlmApiType::OpenAi,
@@ -201,17 +200,13 @@ mod tests {
 
         let input_args_values = vec![text_content.to_string().into()];
 
-        let input_arg_schemas = vec![build_arg_schema(
-            "text",
-            text_content.to_string().into(),
-            BasicValueType::Str,
-        )];
+        let input_arg_schemas = vec![build_arg_schema("text", BasicValueType::Str)];
 
         let result = test_flow_function(factory, spec, input_arg_schemas, input_args_values).await;
 
         if result.is_err() {
             eprintln!(
-                "test_extract_by_llm_with_util: test_flow_function returned error (potentially expected for evaluate): {:?}",
+                "test_extract_by_llm: test_flow_function returned error (potentially expected for evaluate): {:?}",
                 result.as_ref().err()
             );
         }

src/ops/functions/parse_json.rs

@@ -120,12 +120,8 @@ mod tests {
         let input_args_values = vec![json_string_content.to_string().into(), lang_value.clone()];
 
         let input_arg_schemas = vec![
-            build_arg_schema(
-                "text",
-                json_string_content.to_string().into(),
-                BasicValueType::Str,
-            ),
-            build_arg_schema("language", lang_value, BasicValueType::Str),
+            build_arg_schema("text", BasicValueType::Str),
+            build_arg_schema("language", BasicValueType::Str),
         ];
 
         let result = test_flow_function(factory, spec, input_arg_schemas, input_args_values).await;

src/ops/functions/split_recursively.rs

@@ -1090,11 +1090,11 @@ mod tests {
         ];
 
         let input_arg_schemas = vec![
-            build_arg_schema("text", text_content.to_string().into(), BasicValueType::Str),
-            build_arg_schema("chunk_size", (15i64).into(), BasicValueType::Int64),
-            build_arg_schema("min_chunk_size", (5i64).into(), BasicValueType::Int64),
-            build_arg_schema("chunk_overlap", (0i64).into(), BasicValueType::Int64),
-            build_arg_schema("language", Value::Null, BasicValueType::Str),
+            build_arg_schema("text", BasicValueType::Str),
+            build_arg_schema("chunk_size", BasicValueType::Int64),
+            build_arg_schema("min_chunk_size", BasicValueType::Int64),
+            build_arg_schema("chunk_overlap", BasicValueType::Int64),
+            build_arg_schema("language", BasicValueType::Str),
         ];
 
         let result = test_flow_function(factory, spec, input_arg_schemas, input_args_values).await;

src/ops/functions/test_utils.rs

@@ -1,4 +1,6 @@
-use crate::builder::plan::AnalyzedValueMapping;
+use crate::builder::plan::{
+    AnalyzedFieldReference, AnalyzedLocalFieldReference, AnalyzedValueMapping,
+};
 use crate::ops::sdk::{
     AuthRegistry, BasicValueType, EnrichedValueType, FlowInstanceContext, OpArgSchema,
     OpArgsResolver, SimpleFunctionExecutor, SimpleFunctionFactoryBase, Value, make_output_type,
@@ -7,47 +9,52 @@ use anyhow::Result;
 use serde::de::DeserializeOwned;
 use std::sync::Arc;
 
-fn new_literal_op_arg_schema(
-    name: Option<&str>,
-    value: Value,
-    value_type: EnrichedValueType,
-) -> OpArgSchema {
-    OpArgSchema {
-        name: name.map_or(crate::base::spec::OpArgName(None), |n| {
-            crate::base::spec::OpArgName(Some(n.to_string()))
-        }),
-        value_type,
-        analyzed_value: AnalyzedValueMapping::Constant { value },
-    }
-}
-
-// This function provides a helper to create OpArgSchema for literal values.
-pub fn build_arg_schema(name: &str, value: Value, value_type: BasicValueType) -> OpArgSchema {
-    new_literal_op_arg_schema(Some(name), value, make_output_type(value_type))
+// This function builds an argument schema for a flow function.
+pub fn build_arg_schema(
+    name: &str,
+    value_type: BasicValueType,
+) -> (Option<&str>, EnrichedValueType) {
+    (Some(name), make_output_type(value_type))
 }
 
 // This function tests a flow function by providing a spec, input argument schemas, and values.
 pub async fn test_flow_function<S, R, F>(
     factory: Arc<F>,
     spec: S,
-    input_arg_schemas: Vec<OpArgSchema>,
+    input_arg_schemas: Vec<(Option<&str>, EnrichedValueType)>,
     input_arg_values: Vec<Value>,
 ) -> Result<Value>
 where
     S: DeserializeOwned + Send + Sync + 'static,
     R: Send + Sync + 'static,
     F: SimpleFunctionFactoryBase<Spec = S, ResolvedArgs = R> + ?Sized,
 {
+    // 1. Construct OpArgSchema
+    let op_arg_schemas: Vec<OpArgSchema> = input_arg_schemas
+        .into_iter()
+        .enumerate()
+        .map(|(idx, (name, value_type))| OpArgSchema {
+            name: name.map_or(crate::base::spec::OpArgName(None), |n| {
+                crate::base::spec::OpArgName(Some(n.to_string()))
+            }),
+            value_type,
+            analyzed_value: AnalyzedValueMapping::Field(AnalyzedFieldReference {
+                local: AnalyzedLocalFieldReference {
+                    fields_idx: vec![idx as u32],
+                },
+                scope_up_level: 0,
+            }),
+        })
+        .collect();
+
+    // 2. Resolve Schema & Args
+    let mut args_resolver = OpArgsResolver::new(&op_arg_schemas)?;
     let context = Arc::new(FlowInstanceContext {
         flow_instance_name: "test_flow_function".to_string(),
         auth_registry: Arc::new(AuthRegistry::default()),
         py_exec_ctx: None,
     });
 
-    // 2. Resolve Schema & Args
-    // The caller of test_flow_function will be responsible for creating these schemas.
-    let mut args_resolver = OpArgsResolver::new(&input_arg_schemas)?;
-
     let (resolved_args_from_schema, _output_schema): (R, EnrichedValueType) = factory
         .resolve_schema(&spec, &mut args_resolver, &context)
         .await?;