Skip to main content

rig_typesafeai/
wire.rs

1//! Typesafe's unary evaluation operation, using Rig's shared HTTP driver.
2
3use 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
15/// A batch of independent questions about one state.
16pub struct Evaluation;
17
18impl Operation for Evaluation {
19    type Request = Request;
20    // No events: the reply's end is the response.
21    type Event = std::convert::Infallible;
22    type End = Response;
23    type Response = Response;
24    type Fold = Whole<Self>;
25    type Emit = Free;
26
27    /// The transport request id and the reply document reach the response.
28    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/// The settings of Jev's `POST /v1/systemone` endpoint, which are also its
39/// wire. [`connect`](Self::connect) puts them on a transport as a [`Jev`]
40/// client.
41///
42/// Credentials are redacted when this configuration is printed or serialized.
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct JevConfig {
45    token: Secret,
46    model: String,
47    endpoint: String,
48}
49
50impl JevConfig {
51    /// Use `jev-latest` at the Typesafe API.
52    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    /// Read the credential from `JEV_TOKEN` without loading a secrets file.
61    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    /// Select the provider model identifier.
73    pub fn model(mut self, model: impl Into<String>) -> Self {
74        self.model = model.into();
75        self
76    }
77
78    /// Set the complete endpoint URL, including `/v1/systemone`.
79    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
80        self.endpoint = endpoint.into();
81        self
82    }
83
84    /// A client that sends these settings' requests through `http`.
85    pub fn connect(self, http: impl HttpClientExt + 'static) -> Jev {
86        Jev {
87            config: self,
88            http: DynHttpClient::new(http),
89        }
90    }
91
92    /// A client that sends through the shared reqwest client, built once
93    /// per process.
94    #[cfg(feature = "reqwest")]
95    pub fn client(self) -> Jev {
96        self.connect(rig_reqwest::shared())
97    }
98}
99
100/// Jev: its [`JevConfig`] on a transport. The model it builds sends through
101/// that transport.
102#[derive(Debug, Clone)]
103pub struct Jev {
104    config: JevConfig,
105    http: DynHttpClient,
106}
107
108impl Jev {
109    /// `jev-latest` at the Typesafe API with `token`, on the shared reqwest
110    /// client.
111    #[cfg(feature = "reqwest")]
112    pub fn new(token: impl Into<String>) -> Self {
113        JevConfig::new(token).client()
114    }
115
116    /// The credential from `JEV_TOKEN`, on the shared reqwest client.
117    #[cfg(feature = "reqwest")]
118    pub fn from_env() -> Result<Self, EnvError> {
119        Ok(JevConfig::from_env()?.client())
120    }
121
122    /// The settings this client sends with.
123    pub fn config(&self) -> &JevConfig {
124        &self.config
125    }
126
127    /// The evaluation model: the configured Jev model on this client's
128    /// transport.
129    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    // One JSON document in, the reply's end out, the document kept as `raw`.
140    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}