starweaver_runtime/output/
function.rs1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use starweaver_model::ToolDefinition;
5
6use crate::run::AgentRunState;
7
8use super::{validation::parse_output, OutputSchema, OutputValidationResult, OutputValue};
9
10#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
12pub struct OutputFunctionDefinition {
13 pub name: String,
15 #[serde(default, skip_serializing_if = "Option::is_none")]
17 pub description: Option<String>,
18 pub parameters: Value,
20}
21
22impl OutputFunctionDefinition {
23 #[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 #[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 #[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 #[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#[derive(Clone)]
65pub struct SchemaOutputFunction {
66 schema: OutputSchema,
67}
68
69impl SchemaOutputFunction {
70 #[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#[derive(Clone, Debug)]
94pub struct OutputFunctionContext {
95 pub state: AgentRunState,
97}
98
99#[async_trait]
101pub trait OutputFunction: Send + Sync {
102 fn definition(&self) -> OutputFunctionDefinition;
104
105 async fn call(
111 &self,
112 context: OutputFunctionContext,
113 arguments: Value,
114 ) -> OutputValidationResult<OutputValue>;
115}
116
117pub struct FunctionOutputFunction<F> {
119 definition: OutputFunctionDefinition,
120 function: F,
121}
122
123impl<F> FunctionOutputFunction<F> {
124 #[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}