50 lines
1.4 KiB
Rust
50 lines
1.4 KiB
Rust
use async_openai::types::chat::{ChatCompletionTool, ChatCompletionTools, FunctionObjectArgs};
|
|
use schemars::{JsonSchema, Schema, schema_for};
|
|
use serde::de::DeserializeOwned;
|
|
|
|
use crate::artifacts::Artifact;
|
|
|
|
pub trait DataSource: ToolDescription {
|
|
type Args: JsonSchema + DeserializeOwned;
|
|
type Error;
|
|
fn synchronize(&self, artifact: &mut Artifact) -> impl Future<Output = Result<Vec<Artifact>, Self::Error>>;
|
|
fn query(&self, args: &Self::Args) -> impl Future<Output = Result<Vec<Artifact>, Self::Error>>;
|
|
}
|
|
|
|
pub trait ToolDescription {
|
|
fn description(&self) -> &str;
|
|
fn name(&self) -> &str;
|
|
}
|
|
|
|
pub struct Tool {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub schema: Schema
|
|
}
|
|
|
|
impl Tool {
|
|
pub fn from_datasource<T: DataSource>(src: &T) -> Self {
|
|
Self {
|
|
name: src.name().to_string(),
|
|
description: src.description().to_string(),
|
|
schema: schema_for!(T::Args)
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<Tool> for ChatCompletionTool {
|
|
fn from(val: Tool) -> Self {
|
|
ChatCompletionTool {
|
|
function: FunctionObjectArgs::default()
|
|
.name(val.name)
|
|
.description(val.description)
|
|
.parameters(val.schema).build().unwrap()
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<Tool> for ChatCompletionTools {
|
|
fn from(val: Tool) -> Self {
|
|
ChatCompletionTools::Function(val.into())
|
|
}
|
|
} |