1use rig_core::client::{EnvError, env};
4use rig_core::driver::Model;
5use rig_core::error::EncodeError;
6use rig_core::http_client::{DynHttpClient, HttpClientExt};
7use rig_core::operation::Whole;
8use rig_core::wire::{
9 Body, Call, Descriptor, Encoded, Framing, Free, Json, Mode, Operation, Secret, Wire,
10};
11use serde::{Deserialize, Serialize};
12
13use crate::types::{Request, Response};
14
15pub struct Evaluation;
17
18impl Operation for Evaluation {
19 type Request = Request;
20 type Event = std::convert::Infallible;
22 type End = Response;
23 type Response = Response;
24 type Fold = Whole<Self>;
25 type Emit = Free;
26
27 fn fold(_request: &Request, _call: &mut Call<'_>) -> Whole<Self> {
29 Whole::<Self>::stamping(|response, reply| {
30 response
31 .provider_request_id
32 .clone_from(&reply.provider_request_id);
33 response.raw.clone_from(&reply.raw);
34 })
35 }
36}
37
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct JevConfig {
45 token: Secret,
46 model: String,
47 endpoint: String,
48}
49
50impl JevConfig {
51 pub fn new(token: impl Into<String>) -> Self {
53 Self {
54 token: token.into().into(),
55 model: "jev-latest".into(),
56 endpoint: "https://api.typesafe.ai/v1/systemone".into(),
57 }
58 }
59
60 pub fn from_env() -> Result<Self, EnvError> {
62 let token = env::required("JEV_TOKEN")?;
63 if token.trim().is_empty() {
64 return Err(EnvError::Invalid {
65 name: "JEV_TOKEN",
66 detail: "credential is empty".into(),
67 });
68 }
69 Ok(Self::new(token))
70 }
71
72 pub fn model(mut self, model: impl Into<String>) -> Self {
74 self.model = model.into();
75 self
76 }
77
78 pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
80 self.endpoint = endpoint.into();
81 self
82 }
83
84 pub fn connect(self, http: impl HttpClientExt + 'static) -> Jev {
86 Jev {
87 config: self,
88 http: DynHttpClient::new(http),
89 }
90 }
91
92 #[cfg(feature = "reqwest")]
95 pub fn client(self) -> Jev {
96 self.connect(rig_reqwest::shared())
97 }
98}
99
100#[derive(Debug, Clone)]
103pub struct Jev {
104 config: JevConfig,
105 http: DynHttpClient,
106}
107
108impl Jev {
109 #[cfg(feature = "reqwest")]
112 pub fn new(token: impl Into<String>) -> Self {
113 JevConfig::new(token).client()
114 }
115
116 #[cfg(feature = "reqwest")]
118 pub fn from_env() -> Result<Self, EnvError> {
119 Ok(JevConfig::from_env()?.client())
120 }
121
122 pub fn config(&self) -> &JevConfig {
124 &self.config
125 }
126
127 pub fn evaluation(&self) -> Model<JevConfig> {
130 Model::new(self.config.clone(), self.http.clone())
131 }
132}
133
134impl Wire for JevConfig {
135 type Op = Evaluation;
136 type Payload = Encoded;
137 type Frame = rig_core::wire::WireFrame;
138 type Event = (Response, serde_json::Value);
139 type Decoder<'id> = Json;
141 type Reassembler = rig_core::wire::document::Unreassembled;
142
143 fn describe(&self) -> Descriptor<'_> {
144 Descriptor::new("typesafeai").model(self.model.as_str())
145 }
146 fn encode(&self, request: Request, _mode: Mode) -> Result<Encoded, EncodeError> {
147 if self.token.is_empty() || self.model.trim().is_empty() {
148 return Err(EncodeError::request(
149 "credential and model must be nonempty",
150 ));
151 }
152 #[derive(Serialize)]
153 struct Payload<'a> {
154 model: &'a str,
155 #[serde(flatten)]
156 request: Request,
157 }
158 let body = Payload {
159 model: &self.model,
160 request,
161 };
162 let mut authorization =
163 http::HeaderValue::from_str(&format!("Bearer {}", self.token.expose()))
164 .map_err(EncodeError::request)?;
165 authorization.set_sensitive(true);
166 let request = http::Request::post(&self.endpoint)
167 .header(http::header::CONTENT_TYPE, "application/json")
168 .header(http::header::AUTHORIZATION, authorization)
169 .body(Body::Bytes(serde_json::to_vec(&body)?))?;
170 Ok(Encoded::new(request, Framing::Whole)
171 .with_request_id_header(Some("x-typesafe-request-id"))
172 .with_route(Some("/v1/systemone")))
173 }
174
175 fn decoder<'id>(&self) -> Self::Decoder<'id> {
176 Json
177 }
178}