asterai 1.0.0-alpha.16

CLI for asterai - the portable AI compute platform
use crate::command::env::EnvArgs;
use crate::local_store::LocalStore;
use crate::runtime::build_runtime;
use asterai_runtime::component::function_name::ComponentFunctionName;
use asterai_runtime::component::{ComponentId, PackageName, Version};
use asterai_runtime::runtime::parsing::{ValExt, json_value_to_val, parse_primitive};
use asterai_runtime::runtime::{ComponentRuntime, Val};
use eyre::{OptionExt, bail};
use std::str::FromStr;
use wit_parser::{Resolve, TypeDef, TypeDefKind};

impl EnvArgs {
    pub async fn call(&self) -> eyre::Result<()> {
        let resource_id = self.resource_id()?;
        let comp_arg = self.component_arg.as_ref().unwrap();
        println!(
            "calling env {resource_id}'s {}:{} \
             component function {}",
            comp_arg.resolved_namespace(),
            comp_arg.name(),
            self.function.as_ref().unwrap()
        );
        let result = self.call_returning().await?;
        if let Some(output) = result {
            println!("{output}");
        }
        Ok(())
    }

    /// Call a function and return the output as a string.
    pub async fn call_returning(&self) -> eyre::Result<Option<String>> {
        let resource_id = self.resource_id()?;
        let comp_arg = self.component_arg.as_ref().unwrap();
        let comp_ns = comp_arg.resolved_namespace();
        let comp_name = comp_arg.name();
        let comp_id = ComponentId::from_str(&format!("{comp_ns}:{comp_name}"))?;
        let function_string = self.function.clone().unwrap();
        let environment = LocalStore::fetch_environment(&resource_id)
            .map_err(|_| eyre::eyre!("environment '{}' not found locally", resource_id))?;
        let mut runtime = build_runtime(environment, &self.allow_dirs).await?;
        call_on_runtime(&mut runtime, &comp_id, function_string, &self.function_args).await
    }
}

/// Find a function on a runtime, parse inputs, call it, and return formatted output.
pub(crate) async fn call_on_runtime(
    runtime: &mut ComponentRuntime,
    comp_id: &ComponentId,
    function_string: String,
    function_args: &[String],
) -> eyre::Result<Option<String>> {
    let (function_name, package_name_opt) = parse_function_string_into_parts(function_string)?;
    let function = runtime
        .find_function(comp_id, &function_name, package_name_opt)?
        .ok_or_eyre("function not found")?;
    let resolve = runtime
        .resolve_for(comp_id)
        .ok_or_eyre("component not found")?;
    let inputs = parse_inputs_from_string_args(function_args, &function.inputs, &resolve)?;
    let output_opt = runtime.call_function(function, &inputs).await?;
    if let Some(output) = output_opt
        && let Some(function_output) = output.function_output_opt
    {
        let json = function_output.value.val.try_into_json_value();
        return Ok(json.map(|j| match j {
            serde_json::Value::String(s) => s,
            other => other.to_string(),
        }));
    }
    Ok(None)
}

fn parse_function_string_into_parts(
    function_str: String,
) -> eyre::Result<(ComponentFunctionName, Option<PackageName>)> {
    let (package_name_str, function) = match function_str.split_once('/') {
        None => {
            return Ok((
                ComponentFunctionName {
                    interface: None,
                    name: function_str,
                },
                None,
            ));
        }
        Some((a, b)) => {
            if !a.contains(":") {
                return Ok((
                    ComponentFunctionName {
                        interface: Some(a.to_owned()),
                        name: b.to_owned(),
                    },
                    None,
                ));
            }
            (a, b)
        }
    };
    let package_name = package_name_from_str(package_name_str)?;
    let (function_name, interface_opt) = match function.split_once('/') {
        None => (function, None),
        Some((a, b)) => (a, Some(b)),
    };
    Ok((
        ComponentFunctionName {
            interface: interface_opt.map(|v| v.to_owned()),
            name: function_name.to_owned(),
        },
        Some(package_name),
    ))
}

fn parse_inputs_from_string_args(
    args: &[String],
    expected_inputs: &[(String, TypeDef)],
    resolve: &Resolve,
) -> eyre::Result<Vec<Val>> {
    if args.len() != expected_inputs.len() {
        bail!(
            "expected {} argument(s), got {}",
            expected_inputs.len(),
            args.len()
        );
    }
    args.iter()
        .zip(expected_inputs.iter())
        .map(|(arg, (name, type_def))| {
            parse_arg(arg, type_def, resolve)
                .map_err(|e| eyre::eyre!("failed to parse argument '{name}': {e}"))
        })
        .collect()
}

fn parse_arg(arg: &str, type_def: &TypeDef, resolve: &Resolve) -> eyre::Result<Val> {
    match &type_def.kind {
        TypeDefKind::Type(ty) => parse_primitive(strip_quotes(arg), ty),
        TypeDefKind::Record(record) => {
            let json: serde_json::Value = serde_json::from_str(arg)
                .map_err(|e| eyre::eyre!("expected JSON for record: {e}"))?;
            let serde_json::Value::Object(map) = json else {
                bail!("expected JSON object for record");
            };
            let fields = record
                .fields
                .iter()
                .map(|field| {
                    let value = map
                        .get(&field.name)
                        .ok_or_else(|| eyre::eyre!("missing field '{}'", field.name))?;
                    let val = json_value_to_val(value, &field.ty, resolve)?;
                    Ok((field.name.clone(), val))
                })
                .collect::<eyre::Result<Vec<_>>>()?;
            Ok(Val::Record(fields))
        }
        TypeDefKind::List(ty) => {
            let json: serde_json::Value = serde_json::from_str(arg)
                .map_err(|e| eyre::eyre!("expected JSON for list: {e}"))?;
            let serde_json::Value::Array(arr) = json else {
                bail!("expected JSON array for list");
            };
            let vals = arr
                .iter()
                .map(|v| json_value_to_val(v, ty, resolve))
                .collect::<eyre::Result<Vec<_>>>()?;
            Ok(Val::List(vals))
        }
        TypeDefKind::Tuple(tuple) => {
            let json: serde_json::Value = serde_json::from_str(arg)
                .map_err(|e| eyre::eyre!("expected JSON for tuple: {e}"))?;
            let serde_json::Value::Array(arr) = json else {
                bail!("expected JSON array for tuple");
            };
            if arr.len() != tuple.types.len() {
                bail!(
                    "tuple has {} elements, got {}",
                    tuple.types.len(),
                    arr.len()
                );
            }
            let vals = arr
                .iter()
                .zip(tuple.types.iter())
                .map(|(v, ty)| json_value_to_val(v, ty, resolve))
                .collect::<eyre::Result<Vec<_>>>()?;
            Ok(Val::Tuple(vals))
        }
        TypeDefKind::Enum(e) => {
            let arg = strip_quotes(arg);
            let is_valid = e.cases.iter().any(|c| c.name == arg);
            if !is_valid {
                let cases: Vec<&str> = e.cases.iter().map(|c| c.name.as_str()).collect();
                bail!("invalid enum value '{arg}', expected one of: {cases:?}");
            }
            Ok(Val::Enum(arg.to_owned()))
        }
        TypeDefKind::Option(ty) => {
            let arg_stripped = strip_quotes(arg);
            if arg_stripped == "null" || arg_stripped == "none" {
                return Ok(Val::Option(None));
            }
            let inner = json_value_to_val(&serde_json::from_str(arg)?, ty, resolve)?;
            Ok(Val::Option(Some(Box::new(inner))))
        }
        TypeDefKind::Flags(flags) => {
            let json: serde_json::Value = serde_json::from_str(arg)
                .map_err(|e| eyre::eyre!("expected JSON for flags: {e}"))?;
            let serde_json::Value::Array(arr) = json else {
                bail!("expected JSON array for flags");
            };
            let names = arr
                .iter()
                .map(|v| match v {
                    serde_json::Value::String(s) => {
                        let is_valid = flags.flags.iter().any(|f| f.name == *s);
                        if !is_valid {
                            bail!("invalid flag '{s}'");
                        }
                        Ok(s.clone())
                    }
                    _ => bail!("expected string for flag name"),
                })
                .collect::<eyre::Result<Vec<_>>>()?;
            Ok(Val::Flags(names))
        }
        _ => {
            bail!("unsupported type: {:#?}", type_def.kind);
        }
    }
}

/// Strips surrounding double quotes from a string if present.
fn strip_quotes(s: &str) -> &str {
    s.strip_prefix('"')
        .and_then(|s| s.strip_suffix('"'))
        .unwrap_or(s)
}

fn package_name_from_str(package_name_str: &str) -> eyre::Result<PackageName> {
    let (id, version_str_opt) = match package_name_str.split_once('@') {
        None => (package_name_str, None),
        Some((a, b)) => (a, Some(b)),
    };
    let (namespace, name) = match id.split_once(':') {
        None => bail!("invalid package name"),
        Some((a, b)) => (a, b),
    };
    let version_opt = match version_str_opt {
        Some(v) => Some(Version::from_str(v)?),
        None => None,
    };
    Ok(PackageName {
        namespace: namespace.to_owned(),
        name: name.to_owned(),
        version: version_opt,
    })
}