swiftide/examples/tool_custom_schema.rs

76 lines
2.0 KiB
Rust

use std::borrow::Cow;
use anyhow::Result;
use schemars::{JsonSchema, Schema, schema_for};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use swiftide::chat_completion::{Tool, ToolCall, ToolOutput, ToolSpec, errors::ToolError};
use swiftide::traits::AgentContext;
#[derive(Clone)]
struct WorkflowTool;
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[schemars(
title = "WorkflowInstruction",
description = "Choose a workflow action and optional payload",
deny_unknown_fields
)]
struct WorkflowInstruction {
#[schemars(description = "Which workflow action to execute")]
action: WorkflowAction,
#[schemars(description = "Optional payload forwarded to the workflow engine")]
#[serde(default, skip_serializing_if = "Option::is_none")]
payload: Option<Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "lowercase")]
enum WorkflowAction {
Start,
Stop,
Status,
}
#[swiftide::reexports::async_trait::async_trait]
impl Tool for WorkflowTool {
async fn invoke(
&self,
_agent_context: &dyn AgentContext,
_tool_call: &ToolCall,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput::text(
"Workflow execution not implemented in this example",
))
}
fn name<'tool>(&'tool self) -> Cow<'tool, str> {
Cow::Borrowed("workflow_tool")
}
fn tool_spec(&self) -> ToolSpec {
ToolSpec::builder()
.name("workflow_tool")
.description("Executes a workflow action with strict input choices")
.parameters_schema(workflow_schema())
.build()
.expect("tool spec should be valid")
}
}
fn workflow_schema() -> Schema {
schema_for!(WorkflowInstruction)
}
fn main() -> Result<()> {
let tool = WorkflowTool;
let spec = tool.tool_spec();
println!(
"{}",
serde_json::to_string_pretty(&spec).expect("tool spec should serialize"),
);
Ok(())
}