Skip to main content

starweaver_runtime/output/
function.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use starweaver_model::ToolDefinition;
5
6use crate::run::AgentRunState;
7
8use super::{OutputSchema, OutputValidationResult, OutputValue, validation::parse_output};
9
10/// Function-style final output call definition.
11#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
12pub struct OutputFunctionDefinition {
13    /// Output function name exposed to the model.
14    pub name: String,
15    /// Output function description.
16    #[serde(default, skip_serializing_if = "Option::is_none")]
17    pub description: Option<String>,
18    /// JSON schema for output function arguments.
19    pub parameters: Value,
20}
21
22impl OutputFunctionDefinition {
23    /// Build an output function definition.
24    #[must_use]
25    pub fn new(name: impl Into<String>, parameters: Value) -> Self {
26        Self {
27            name: name.into(),
28            description: None,
29            parameters,
30        }
31    }
32
33    /// Add description.
34    #[must_use]
35    pub fn with_description(mut self, description: impl Into<String>) -> Self {
36        self.description = Some(description.into());
37        self
38    }
39
40    /// Build an output function definition from an output schema.
41    #[must_use]
42    pub fn from_output_schema(schema: &OutputSchema) -> Self {
43        let mut definition = Self::new(schema.name.clone(), schema.schema.clone());
44        definition.description.clone_from(&schema.description);
45        definition
46    }
47
48    /// Convert to a provider-neutral tool definition.
49    #[must_use]
50    pub fn tool_definition(&self) -> ToolDefinition {
51        ToolDefinition {
52            name: self.name.clone(),
53            description: self.description.clone(),
54            parameters: self.parameters.clone(),
55            return_schema: None,
56            strict: None,
57            sequential: None,
58            metadata: serde_json::Map::new(),
59        }
60    }
61}
62
63/// Schema-backed output function used for tool-mode structured output.
64#[derive(Clone)]
65pub struct SchemaOutputFunction {
66    schema: OutputSchema,
67}
68
69impl SchemaOutputFunction {
70    /// Build an output function that validates arguments against an output schema.
71    #[must_use]
72    pub const fn new(schema: OutputSchema) -> Self {
73        Self { schema }
74    }
75}
76
77#[async_trait]
78impl OutputFunction for SchemaOutputFunction {
79    fn definition(&self) -> OutputFunctionDefinition {
80        OutputFunctionDefinition::from_output_schema(&self.schema)
81    }
82
83    async fn call(
84        &self,
85        _context: OutputFunctionContext,
86        arguments: Value,
87    ) -> OutputValidationResult<OutputValue> {
88        parse_output(&arguments.to_string(), Some(&self.schema))
89    }
90}
91
92/// Output function execution context.
93#[derive(Clone, Debug)]
94pub struct OutputFunctionContext {
95    /// Run state at finalization time.
96    pub state: AgentRunState,
97}
98
99/// Output function called by the model to finish a run.
100#[async_trait]
101pub trait OutputFunction: Send + Sync {
102    /// Output function definition exposed to the model.
103    fn definition(&self) -> OutputFunctionDefinition;
104
105    /// Execute the output function with model-provided arguments.
106    ///
107    /// # Errors
108    ///
109    /// Returns an error when the output function rejects the arguments or fails.
110    async fn call(
111        &self,
112        context: OutputFunctionContext,
113        arguments: Value,
114    ) -> OutputValidationResult<OutputValue>;
115}
116
117/// Function-backed output function.
118pub struct FunctionOutputFunction<F> {
119    definition: OutputFunctionDefinition,
120    function: F,
121}
122
123impl<F> FunctionOutputFunction<F> {
124    /// Build a function-backed output function.
125    #[must_use]
126    pub const fn new(definition: OutputFunctionDefinition, function: F) -> Self {
127        Self {
128            definition,
129            function,
130        }
131    }
132}
133
134#[async_trait]
135impl<F, Fut> OutputFunction for FunctionOutputFunction<F>
136where
137    F: Send + Sync + Fn(OutputFunctionContext, Value) -> Fut,
138    Fut: Send + std::future::Future<Output = OutputValidationResult<OutputValue>>,
139{
140    fn definition(&self) -> OutputFunctionDefinition {
141        self.definition.clone()
142    }
143
144    async fn call(
145        &self,
146        context: OutputFunctionContext,
147        arguments: Value,
148    ) -> OutputValidationResult<OutputValue> {
149        (self.function)(context, arguments).await
150    }
151}