1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use anyhow::Result;
use async_trait::async_trait;
use core::fmt::Debug;
use schemars::JsonSchema;
use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;

#[async_trait]
pub trait Cog: Sized + Send {
	type Request: DeserializeOwned + JsonSchema;
	type Response: CogResponse + Debug + JsonSchema;

	/// Setup the cog
	///
	/// # Errors
	///
	/// Returns an error if setup fails.
	async fn setup() -> Result<Self>;

	/// Run a prediction
	async fn predict(&self, input: Self::Request) -> Result<Self::Response>;
}

/// A response from a cog
pub trait CogResponse: Send {
	fn into_response(self) -> Value;
}

impl<T: Serialize + Send> CogResponse for T {
	fn into_response(self) -> Value {
		serde_json::to_value(self).unwrap()
	}
}