fx-mistral 0.0.2

A library to leverage the Mistral API for OCR and data extraction from PDFs.
Documentation 
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

use serde::{Deserialize, Serialize};

//
// Chat Request structs
//

#[derive(Serialize, Deserialize, Debug)]
pub struct ChatRequest {
    pub model: String,
    pub messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ResponseFormat {
    #[serde(rename = "type")]
    pub format_type: String, // should be "json_schema"
    pub json_schema: JsonSchemaFormat,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct JsonSchemaFormat {
    pub schema: serde_json::Value, // Accept raw JSON schema as a serde_json::Value
    pub name: String,
    pub strict: bool,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub enum Message {
    Simple {
        role: String,
        content: String, // For system messages
    },
    WithContentArray {
        role: String,
        content: Vec<Content>, // For user messages with docs/texts
    },
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "type")]
pub enum Content {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "document_url")]
    DocumentUrl { document_url: String },
}

pub struct Messages {
    pub messages: Vec<Message>,
}

//
// Chat Request Builder
//

pub struct ChatRequestBuilder {
    model: String,
    messages: Vec<Message>,
    response_format: Option<ResponseFormat>,
    max_tokens: Option<u32>,
    temperature: Option<f32>,
}

impl ChatRequestBuilder {
    pub fn new<S: Into<String>>(model: S, system_message: S, temparatur: f32) -> Self {
        ChatRequestBuilder {
            model: model.into(),
            messages: vec![Message::Simple {
                role: "system".to_string(),
                content: system_message.into(),
            }],
            response_format: None,
            max_tokens: None,
            temperature: Some(temparatur),
        }
    }

    pub fn add_user_message<S: Into<String>>(mut self, text: S) -> Self {
        self.messages.push(Message::WithContentArray {
            role: "user".to_string(),
            content: vec![Content::Text { text: text.into() }],
        });
        self
    }

    pub fn add_document_message<S: Into<String>>(mut self, text: S, document_url: S) -> Self {
        self.messages.push(Message::WithContentArray {
            role: "user".to_string(),
            content: vec![
                Content::Text { text: text.into() },
                Content::DocumentUrl {
                    document_url: document_url.into(),
                },
            ],
        });
        self
    }

    pub fn response_format_from_json<S: Into<String>>(mut self, schema_json: S, name: S, strict: bool) -> Self {
        let schema_value: serde_json::Value = serde_json::from_str(&schema_json.into()).expect("Invalid JSON schema");
        self.response_format = Some(ResponseFormat {
            format_type: "json_schema".to_string(),
            json_schema: JsonSchemaFormat {
                schema: schema_value,
                name: name.into(),
                strict,
            },
        });
        self
    }

    pub fn max_tokens(mut self, value: u32) -> Self {
        self.max_tokens = Some(value);
        self
    }

    pub fn temperature(mut self, value: f32) -> Self {
        self.temperature = Some(value);
        self
    }

    pub fn build(self) -> ChatRequest {
        ChatRequest {
            model: self.model,
            messages: self.messages,
            response_format: self.response_format,
            max_tokens: self.max_tokens,
            temperature: self.temperature,
        }
    }
}