scrapfly_sdk/config/
extraction.rs

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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
//! Extraction endpoint configuration — ported from `sdk/go/config_extraction.go`.

use crate::enums::{CompressionFormat, ExtractionModel};
use crate::error::ScrapflyError;

use super::url_safe_b64_encode;

/// Configuration for a `POST /extraction` request.
#[derive(Debug, Clone, Default)]
pub struct ExtractionConfig {
    /// Document bytes (required).
    pub body: Vec<u8>,
    /// Content type, e.g. `text/html` (required).
    pub content_type: String,
    /// Original URL (helps the AI with context).
    pub url: Option<String>,
    /// Character set.
    pub charset: Option<String>,
    /// Saved extraction template name.
    pub extraction_template: Option<String>,
    /// Inline (ephemeral) template.
    pub extraction_ephemeral_template: Option<serde_json::Value>,
    /// AI extraction prompt.
    pub extraction_prompt: Option<String>,
    /// Extraction model.
    pub extraction_model: Option<ExtractionModel>,
    /// Body is compressed.
    pub is_document_compressed: bool,
    /// Compression format.
    pub document_compression_format: Option<CompressionFormat>,
    /// Webhook name.
    pub webhook: Option<String>,
    /// Maximum time in seconds for extraction processing.
    pub timeout: Option<u32>,
}

impl ExtractionConfig {
    /// Start a builder.
    pub fn builder(body: Vec<u8>, content_type: impl Into<String>) -> ExtractionConfigBuilder {
        ExtractionConfigBuilder {
            cfg: ExtractionConfig {
                body,
                content_type: content_type.into(),
                ..Default::default()
            },
        }
    }

    /// Query params (key is added separately by the client).
    pub fn to_query_pairs(&self) -> Result<Vec<(String, String)>, ScrapflyError> {
        if self.body.is_empty() {
            return Err(ScrapflyError::Config("body is required".into()));
        }
        if self.content_type.is_empty() {
            return Err(ScrapflyError::Config("content_type is required".into()));
        }
        let tpl_count = [
            self.extraction_template.is_some(),
            self.extraction_ephemeral_template.is_some(),
        ]
        .iter()
        .filter(|x| **x)
        .count();
        if tpl_count > 1 {
            return Err(ScrapflyError::Config(
                "cannot use both extraction_template and extraction_ephemeral_template".into(),
            ));
        }

        let mut out = Vec::new();
        out.push(("content_type".into(), self.content_type.clone()));
        if let Some(u) = &self.url {
            out.push(("url".into(), u.clone()));
        }
        if let Some(c) = &self.charset {
            out.push(("charset".into(), c.clone()));
        }
        if let Some(t) = &self.extraction_template {
            out.push(("extraction_template".into(), t.clone()));
        }
        if let Some(t) = &self.extraction_ephemeral_template {
            let s = serde_json::to_string(t)?;
            out.push((
                "extraction_template".into(),
                format!("ephemeral:{}", url_safe_b64_encode(&s)),
            ));
        }
        if let Some(p) = &self.extraction_prompt {
            out.push(("extraction_prompt".into(), p.clone()));
        }
        if let Some(m) = self.extraction_model {
            out.push(("extraction_model".into(), m.as_str().into()));
        }
        if let Some(wh) = &self.webhook {
            out.push(("webhook_name".into(), wh.clone()));
        }
        if let Some(t) = self.timeout {
            out.push(("timeout".into(), t.to_string()));
        }
        Ok(out)
    }
}

/// Builder for [`ExtractionConfig`].
#[derive(Debug, Clone)]
pub struct ExtractionConfigBuilder {
    cfg: ExtractionConfig,
}

impl ExtractionConfigBuilder {
    /// Original URL.
    pub fn url(mut self, v: impl Into<String>) -> Self {
        self.cfg.url = Some(v.into());
        self
    }
    /// Character set.
    pub fn charset(mut self, v: impl Into<String>) -> Self {
        self.cfg.charset = Some(v.into());
        self
    }
    /// Saved template name.
    pub fn extraction_template(mut self, v: impl Into<String>) -> Self {
        self.cfg.extraction_template = Some(v.into());
        self
    }
    /// Inline template.
    pub fn extraction_ephemeral_template(mut self, v: serde_json::Value) -> Self {
        self.cfg.extraction_ephemeral_template = Some(v);
        self
    }
    /// AI prompt.
    pub fn extraction_prompt(mut self, v: impl Into<String>) -> Self {
        self.cfg.extraction_prompt = Some(v.into());
        self
    }
    /// Model.
    pub fn extraction_model(mut self, v: ExtractionModel) -> Self {
        self.cfg.extraction_model = Some(v);
        self
    }
    /// Body is compressed.
    pub fn is_document_compressed(mut self, v: bool) -> Self {
        self.cfg.is_document_compressed = v;
        self
    }
    /// Compression format.
    pub fn document_compression_format(mut self, v: CompressionFormat) -> Self {
        self.cfg.document_compression_format = Some(v);
        self
    }
    /// Webhook name.
    pub fn timeout(mut self, v: u32) -> Self {
        self.cfg.timeout = Some(v);
        self
    }
    pub fn webhook(mut self, v: impl Into<String>) -> Self {
        self.cfg.webhook = Some(v.into());
        self
    }
    /// Finalize the builder.
    pub fn build(self) -> Result<ExtractionConfig, ScrapflyError> {
        if self.cfg.body.is_empty() {
            return Err(ScrapflyError::Config("body is required".into()));
        }
        if self.cfg.content_type.is_empty() {
            return Err(ScrapflyError::Config("content_type is required".into()));
        }
        Ok(self.cfg)
    }
}