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
use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// The `Document` struct represents a document with content, metadata, and a score.
/// The `page_content` field is a string that contains the content of the document.
/// The `metadata` field is a `HashMap` where the keys represent metadata properties and the values represent property values.
/// The `score` field represents a relevance score for the document and is a floating point number.
///
/// # Usage
/// ```rust,ignore
/// let my_doc = Document::new("This is the document content.".to_string())
///    .with_metadata({
///       let mut metadata = HashMap::new();
///       metadata.insert("author".to_string(), json!("John Doe"));
///       metadata
///   })
///    .with_score(0.75);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Document {
    pub page_content: String,
    pub metadata: HashMap<String, Value>,
    pub score: f64,
}

impl Document {
    /// Constructs a new `Document` with provided `page_content`, an empty `metadata` map and a `score` of 0.
    pub fn new<S: Into<String>>(page_content: S) -> Self {
        Document {
            page_content: page_content.into(),
            metadata: HashMap::new(),
            score: 0.0,
        }
    }

    /// Sets the `metadata` Map of the `Document` to the provided HashMap.
    pub fn with_metadata(mut self, metadata: HashMap<String, Value>) -> Self {
        self.metadata = metadata;
        self
    }

    /// Sets the `score` of the `Document` to the provided float.
    pub fn with_score(mut self, score: f64) -> Self {
        self.score = score;
        self
    }
}

impl Default for Document {
    /// Provides a default `Document` with an empty `page_content`, an empty `metadata` map and a `score` of 0.
    fn default() -> Self {
        Document {
            page_content: "".to_string(),
            metadata: HashMap::new(),
            score: 0.0,
        }
    }
}