1use crate::PredictEditsRequestTrigger;
 2use serde::{Deserialize, Serialize};
 3use std::borrow::Cow;
 4use std::ops::Range;
 5use strum::{AsRefStr, EnumString};
 6
 7pub const PREDICT_EDITS_MODE_HEADER_NAME: &str = "X-Zed-Predict-Edits-Mode";
 8
 9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, AsRefStr, EnumString)]
10#[serde(rename_all = "snake_case")]
11#[strum(serialize_all = "snake_case")]
12pub enum PredictEditsMode {
13    Eager,
14    Subtle,
15}
16
17#[derive(Debug, Deserialize, Serialize)]
18pub struct RawCompletionRequest {
19    pub model: String,
20    pub prompt: String,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub max_tokens: Option<u32>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub temperature: Option<f32>,
25    pub stop: Vec<Cow<'static, str>>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub environment: Option<String>,
28}
29
30#[derive(Debug, Serialize, Deserialize)]
31pub struct PredictEditsV3Request {
32    #[serde(flatten)]
33    pub input: zeta_prompt::ZetaPromptInput,
34    #[serde(default)]
35    pub trigger: PredictEditsRequestTrigger,
36}
37
38#[derive(Debug, Deserialize, Serialize)]
39pub struct PredictEditsV3Response {
40    pub request_id: String,
41    pub output: String,
42    /// The editable region byte range within `cursor_excerpt` that the
43    /// server used for this request. When present, the client should use
44    /// this range to extract the old text from its local excerpt for
45    /// diffing, rather than relying on its own format-derived range.
46    pub editable_range: Range<usize>,
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub model_version: Option<String>,
49}
50
51#[derive(Debug, Deserialize, Serialize)]
52pub struct RawCompletionResponse {
53    pub id: String,
54    pub object: String,
55    pub created: u64,
56    pub model: String,
57    pub choices: Vec<RawCompletionChoice>,
58    pub usage: RawCompletionUsage,
59}
60
61#[derive(Debug, Deserialize, Serialize)]
62pub struct RawCompletionChoice {
63    pub text: String,
64    pub finish_reason: Option<String>,
65}
66
67#[derive(Debug, Serialize, Deserialize)]
68pub struct RawCompletionUsage {
69    pub prompt_tokens: u32,
70    pub completion_tokens: u32,
71    pub total_tokens: u32,
72}