use rig::{completion::ToolDefinition, tool::Tool};
use serde::{Deserialize, Serialize};
use serde_json::json;
use crate::{
error::SenseiError,
lua,
rcon_ext::{execute_lua_json, SharedRcon},
};
pub struct GetRecipe {
pub(crate) rcon: SharedRcon,
}
impl GetRecipe {
pub const fn new(rcon: SharedRcon) -> Self {
Self { rcon }
}
}
#[derive(Debug, Deserialize)]
pub struct GetRecipeArgs {
pub recipe_name: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
pub struct RecipeIngredient {
pub name: String,
#[serde(rename = "type")]
pub kind: String,
pub amount: f64,
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
pub struct RecipeProduct {
pub name: String,
#[serde(rename = "type")]
pub kind: String,
pub amount: f64,
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
pub struct RecipeInfo {
pub name: String,
pub energy: f64,
pub ingredients: Vec<RecipeIngredient>,
pub products: Vec<RecipeProduct>,
}
impl Tool for GetRecipe {
const NAME: &'static str = "get_recipe";
type Error = SenseiError;
type Args = GetRecipeArgs;
type Output = RecipeInfo;
async fn definition(&self, _prompt: String) -> ToolDefinition {
ToolDefinition {
name: "get_recipe".to_string(),
description:
"Look up a recipe's ingredients, products, and crafting time by prototype name"
.to_string(),
parameters: json!({
"type": "object",
"properties": {
"recipe_name": {
"type": "string",
"description": "The recipe prototype name (e.g. 'iron-gear-wheel', 'electronic-circuit')"
}
},
"required": ["recipe_name"]
}),
}
}
async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
let lua = lua::recipe(&args.recipe_name);
let json = execute_lua_json(&self.rcon, &lua).await?;
Ok(serde_json::from_str(&json)?)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_recipe() {
let json = r#"{
"name": "iron-gear-wheel",
"energy": 0.5,
"ingredients": [{"name": "iron-plate", "type": "item", "amount": 2}],
"products": [{"name": "iron-gear-wheel", "type": "item", "amount": 1}]
}"#;
let recipe: RecipeInfo = serde_json::from_str(json).unwrap();
assert_eq!(recipe.name, "iron-gear-wheel");
assert_eq!(recipe.energy, 0.5);
assert_eq!(recipe.ingredients.len(), 1);
assert_eq!(recipe.ingredients[0].name, "iron-plate");
assert_eq!(recipe.ingredients[0].amount, 2.0);
assert_eq!(recipe.products[0].name, "iron-gear-wheel");
}
#[test]
fn test_parse_multi_ingredient_recipe() {
let json = r#"{
"name": "electronic-circuit",
"energy": 0.5,
"ingredients": [
{"name": "iron-plate", "type": "item", "amount": 1},
{"name": "copper-cable", "type": "item", "amount": 3}
],
"products": [{"name": "electronic-circuit", "type": "item", "amount": 1}]
}"#;
let recipe: RecipeInfo = serde_json::from_str(json).unwrap();
assert_eq!(recipe.ingredients.len(), 2);
}
}