Expose eval functionality by the CLI.

Author: badmonster0

python/cocoindex/cli.py

@@ -1,4 +1,5 @@
 import click
+import datetime
 
 from . import flow, lib
 from .setup import check_setup_status, CheckSetupStatusOptions, apply_setup_changes
@@ -52,6 +53,24 @@ def update(flow_name: str | None):
     stats = _flow_by_name(flow_name).update()
     print(stats)
 
+@cli.command()
+@click.argument("flow_name", type=str, required=False)
+@click.option(
+    "-o", "--output-dir", type=str, required=False,
+    help="The directory to dump the evaluation output to.")
+@click.option(
+    "-c", "--use-cache", is_flag=True, show_default=True, default=True,
+    help="Use cached evaluation results if available.")
+def evaluate(flow_name: str | None, output_dir: str | None, use_cache: bool = True):
+    """
+    Evaluate and dump the flow.
+    """
+    fl = _flow_by_name(flow_name)
+    if output_dir is None:
+        output_dir = f"eval_{fl.name}_{datetime.datetime.now().strftime('%y%m%d_%H%M%S')}"
+    options = flow.EvaluateAndDumpOptions(output_dir=output_dir, use_cache=use_cache)
+    fl.evaluate_and_dump(options)
+
 _default_server_settings = lib.ServerSettings.from_env()
 
 @cli.command()

python/cocoindex/flow.py

@@ -9,6 +9,7 @@
 from typing import Any, Callable, Sequence, TypeVar, get_origin
 from threading import Lock
 from enum import Enum
+from dataclasses import dataclass
 
 from . import _engine
 from . import vector
@@ -61,18 +62,18 @@ def _create_data_slice(
 def _spec_kind(spec: Any) -> str:
     return spec.__class__.__name__
 
-def _spec_value_dump(v: Any) -> Any:
-    """Recursively dump a spec object and its nested attributes to a dictionary."""
+def _dump_engine_object(v: Any) -> Any:
+    """Recursively dump an object for engine. Engine side uses `Pythonzized` to catch."""
     if isinstance(v, type) or get_origin(v) is not None:
         return encode_enriched_type(v)
     elif isinstance(v, Enum):
         return v.value
     elif hasattr(v, '__dict__'):
-        return {k: _spec_value_dump(v) for k, v in v.__dict__.items()}
+        return {k: _dump_engine_object(v) for k, v in v.__dict__.items()}
     elif isinstance(v, (list, tuple)):
-        return [_spec_value_dump(item) for item in v]
+        return [_dump_engine_object(item) for item in v]
     elif isinstance(v, dict):
-        return {k: _spec_value_dump(v) for k, v in v.items()}
+        return {k: _dump_engine_object(v) for k, v in v.items()}
     return v
 
 T = TypeVar('T')
@@ -177,7 +178,7 @@ def transform(self, fn_spec: op.FunctionSpec, *args, **kwargs) -> DataSlice:
             lambda target_scope, name:
                 flow_builder_state.engine_flow_builder.transform(
                     _spec_kind(fn_spec),
-                    _spec_value_dump(fn_spec),
+                    _dump_engine_object(fn_spec),
                     transform_args,
                     target_scope,
                     flow_builder_state.field_name_builder.build_name(
@@ -267,7 +268,7 @@ def export(self, name: str, target_spec: op.StorageSpec, /, *,
             {"field_name": field_name, "metric": metric.value}
             for field_name, metric in vector_index]
         self._flow_builder_state.engine_flow_builder.export(
-            name, _spec_kind(target_spec), _spec_value_dump(target_spec),
+            name, _spec_kind(target_spec), _dump_engine_object(target_spec),
             index_options, self._engine_data_collector)
 
 
@@ -316,13 +317,20 @@ def add_source(self, spec: op.SourceSpec, /, name: str | None = None) -> DataSli
             self._state,
             lambda target_scope, name: self._state.engine_flow_builder.add_source(
                 _spec_kind(spec),
-                _spec_value_dump(spec),
+                _dump_engine_object(spec),
                 target_scope,
                 self._state.field_name_builder.build_name(
                     name, prefix=_to_snake_case(_spec_kind(spec))+'_'),
             ),
             name
         )
+@dataclass
+class EvaluateAndDumpOptions:
+    """
+    Options for evaluating and dumping a flow.
+    """
+    output_dir: str
+    use_cache: bool = True
 
 class Flow:
     """
@@ -348,20 +356,32 @@ def __str__(self):
     def __repr__(self):
         return repr(self._lazy_engine_flow())
 
+    @property
+    def name(self) -> str:
+        """
+        Get the name of the flow.
+        """
+        return self._lazy_engine_flow().name()
+
     def update(self):
         """
         Update the index defined by the flow.
         Once the function returns, the indice is fresh up to the moment when the function is called.
         """
         return self._lazy_engine_flow().update()
 
+    def evaluate_and_dump(self, options: EvaluateAndDumpOptions):
+        """
+        Evaluate and dump the flow.
+        """
+        return self._lazy_engine_flow().evaluate_and_dump(_dump_engine_object(options))
+
     def internal_flow(self) -> _engine.Flow:
         """
         Get the engine flow.
         """
         return self._lazy_engine_flow()
 
-
 def _create_lazy_flow(name: str | None, fl_def: Callable[[FlowBuilder, DataScope], None]) -> Flow:
     """
     Create a flow without really building it yet.

src/execution/dumper.rs

@@ -15,7 +15,7 @@ use crate::builder::plan::{AnalyzedSourceOp, ExecutionPlan};
 use crate::utils::yaml_ser::YamlSerializer;
 
 #[derive(Debug, Clone, Deserialize)]
-pub struct DumpEvaluationOutputOptions {
+pub struct EvaluateAndDumpOptions {
     pub output_dir: String,
     pub use_cache: bool,
 }
@@ -54,7 +54,7 @@ struct Dumper<'a> {
     plan: &'a ExecutionPlan,
     schema: &'a schema::DataSchema,
     pool: &'a PgPool,
-    options: DumpEvaluationOutputOptions,
+    options: EvaluateAndDumpOptions,
 }
 
 impl<'a> Dumper<'a> {
@@ -215,7 +215,7 @@ impl<'a> Dumper<'a> {
 pub async fn evaluate_and_dump(
     plan: &ExecutionPlan,
     schema: &schema::DataSchema,
-    options: DumpEvaluationOutputOptions,
+    options: EvaluateAndDumpOptions,
     pool: &PgPool,
 ) -> Result<()> {
     let output_dir = Path::new(&options.output_dir);

src/py/mod.rs

@@ -103,6 +103,10 @@ impl Flow {
         self.__str__()
     }
 
+    pub fn name(&self) -> &str {
+        &self.0.flow_instance.name
+    }
+
     pub fn update(&self, py: Python<'_>) -> PyResult<IndexUpdateInfo> {
         py.allow_threads(|| {
             let lib_context = get_lib_context()
@@ -122,7 +126,7 @@ impl Flow {
     pub fn evaluate_and_dump(
         &self,
         py: Python<'_>,
-        options: Pythonized<execution::dumper::DumpEvaluationOutputOptions>,
+        options: Pythonized<execution::dumper::EvaluateAndDumpOptions>,
     ) -> PyResult<()> {
         py.allow_threads(|| {
             let lib_context = get_lib_context()