claude_agent/types/
citations.rs

1//! Citation types for document-based responses.
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
6pub struct CitationsConfig {
7    pub enabled: bool,
8}
9
10impl CitationsConfig {
11    pub const fn enabled() -> Self {
12        Self { enabled: true }
13    }
14
15    pub const fn disabled() -> Self {
16        Self { enabled: false }
17    }
18}
19
20impl Default for CitationsConfig {
21    fn default() -> Self {
22        Self::enabled()
23    }
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
27#[serde(tag = "type", rename_all = "snake_case")]
28pub enum Citation {
29    CharLocation(CharLocationCitation),
30    PageLocation(PageLocationCitation),
31    ContentBlockLocation(ContentBlockLocationCitation),
32    SearchResultLocation(SearchResultLocationCitation),
33    WebSearchResultLocation(WebSearchResultLocationCitation),
34}
35
36impl Citation {
37    pub fn cited_text(&self) -> &str {
38        match self {
39            Self::CharLocation(c) => &c.cited_text,
40            Self::PageLocation(c) => &c.cited_text,
41            Self::ContentBlockLocation(c) => &c.cited_text,
42            Self::SearchResultLocation(c) => &c.cited_text,
43            Self::WebSearchResultLocation(c) => &c.cited_text,
44        }
45    }
46
47    pub fn document_title(&self) -> Option<&str> {
48        match self {
49            Self::CharLocation(c) => c.document_title.as_deref(),
50            Self::PageLocation(c) => c.document_title.as_deref(),
51            Self::ContentBlockLocation(c) => c.document_title.as_deref(),
52            Self::SearchResultLocation(c) => c.title.as_deref(),
53            Self::WebSearchResultLocation(c) => c.title.as_deref(),
54        }
55    }
56
57    pub fn document_index(&self) -> Option<usize> {
58        match self {
59            Self::CharLocation(c) => Some(c.document_index),
60            Self::PageLocation(c) => Some(c.document_index),
61            Self::ContentBlockLocation(c) => Some(c.document_index),
62            Self::SearchResultLocation(_) | Self::WebSearchResultLocation(_) => None,
63        }
64    }
65
66    pub fn is_search_result(&self) -> bool {
67        matches!(
68            self,
69            Self::SearchResultLocation(_) | Self::WebSearchResultLocation(_)
70        )
71    }
72
73    pub fn url(&self) -> Option<&str> {
74        match self {
75            Self::WebSearchResultLocation(c) => Some(&c.url),
76            Self::SearchResultLocation(c) => Some(&c.source),
77            _ => None,
78        }
79    }
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
83pub struct CharLocationCitation {
84    pub cited_text: String,
85    pub document_index: usize,
86    #[serde(skip_serializing_if = "Option::is_none")]
87    pub document_title: Option<String>,
88    pub start_char_index: usize,
89    pub end_char_index: usize,
90}
91
92#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
93pub struct PageLocationCitation {
94    pub cited_text: String,
95    pub document_index: usize,
96    #[serde(skip_serializing_if = "Option::is_none")]
97    pub document_title: Option<String>,
98    pub start_page_number: u32,
99    pub end_page_number: u32,
100}
101
102#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
103pub struct ContentBlockLocationCitation {
104    pub cited_text: String,
105    pub document_index: usize,
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pub document_title: Option<String>,
108    pub start_block_index: usize,
109    pub end_block_index: usize,
110}
111
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
113pub struct SearchResultLocationCitation {
114    pub cited_text: String,
115    pub source: String,
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub title: Option<String>,
118    pub search_result_index: usize,
119    pub start_block_index: usize,
120    pub end_block_index: usize,
121}
122
123#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
124pub struct WebSearchResultLocationCitation {
125    pub cited_text: String,
126    pub url: String,
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub title: Option<String>,
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub encrypted_index: Option<String>,
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn test_citation_serialization() {
139        let citation = Citation::CharLocation(CharLocationCitation {
140            cited_text: "test".to_string(),
141            document_index: 0,
142            document_title: Some("Doc".to_string()),
143            start_char_index: 0,
144            end_char_index: 4,
145        });
146
147        let json = serde_json::to_string(&citation).unwrap();
148        assert!(json.contains("char_location"));
149
150        let parsed: Citation = serde_json::from_str(&json).unwrap();
151        assert_eq!(parsed.cited_text(), "test");
152    }
153
154    #[test]
155    fn test_search_result_citation() {
156        let citation = Citation::SearchResultLocation(SearchResultLocationCitation {
157            cited_text: "search text".to_string(),
158            source: "https://example.com".to_string(),
159            title: Some("Example".to_string()),
160            search_result_index: 0,
161            start_block_index: 0,
162            end_block_index: 1,
163        });
164
165        assert!(citation.is_search_result());
166        assert!(citation.document_index().is_none());
167    }
168}
claude_agent/types/citations.rs

claude_agent/types/
citations.rs
