codex_utils_stream_parser/
citation.rs

1use crate::InlineHiddenTagParser;
2use crate::InlineTagSpec;
3use crate::StreamTextChunk;
4use crate::StreamTextParser;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7enum CitationTag {
8    Citation,
9}
10
11const CITATION_OPEN: &str = "<oai-mem-citation>";
12const CITATION_CLOSE: &str = "</oai-mem-citation>";
13
14/// Stream parser for `<oai-mem-citation>...</oai-mem-citation>` tags.
15///
16/// This is a thin convenience wrapper around [`InlineHiddenTagParser`]. It returns citation bodies
17/// as plain strings and omits the citation tags from visible text.
18///
19/// Matching is literal and non-nested. If EOF is reached before a closing
20/// `</oai-mem-citation>`, the parser auto-closes the tag and returns the buffered body as an
21/// extracted citation.
22#[derive(Debug)]
23pub struct CitationStreamParser {
24    inner: InlineHiddenTagParser<CitationTag>,
25}
26
27impl CitationStreamParser {
28    pub fn new() -> Self {
29        Self {
30            inner: InlineHiddenTagParser::new(vec![InlineTagSpec {
31                tag: CitationTag::Citation,
32                open: CITATION_OPEN,
33                close: CITATION_CLOSE,
34            }]),
35        }
36    }
37}
38
39impl Default for CitationStreamParser {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl StreamTextParser for CitationStreamParser {
46    type Extracted = String;
47
48    fn push_str(&mut self, chunk: &str) -> StreamTextChunk<Self::Extracted> {
49        let inner = self.inner.push_str(chunk);
50        StreamTextChunk {
51            visible_text: inner.visible_text,
52            extracted: inner.extracted.into_iter().map(|tag| tag.content).collect(),
53        }
54    }
55
56    fn finish(&mut self) -> StreamTextChunk<Self::Extracted> {
57        let inner = self.inner.finish();
58        StreamTextChunk {
59            visible_text: inner.visible_text,
60            extracted: inner.extracted.into_iter().map(|tag| tag.content).collect(),
61        }
62    }
63}
64
65/// Strip citation tags from a complete string and return `(visible_text, citations)`.
66///
67/// This uses [`CitationStreamParser`] internally, so it inherits the same semantics:
68/// literal, non-nested matching and auto-closing unterminated citations at EOF.
69pub fn strip_citations(text: &str) -> (String, Vec<String>) {
70    let mut parser = CitationStreamParser::new();
71    let mut out = parser.push_str(text);
72    let tail = parser.finish();
73    out.visible_text.push_str(&tail.visible_text);
74    out.extracted.extend(tail.extracted);
75    (out.visible_text, out.extracted)
76}
77
78#[cfg(test)]
79mod tests {
80    use super::CitationStreamParser;
81    use super::strip_citations;
82    use crate::StreamTextChunk;
83    use crate::StreamTextParser;
84    use pretty_assertions::assert_eq;
85
86    fn collect_chunks<P>(parser: &mut P, chunks: &[&str]) -> StreamTextChunk<P::Extracted>
87    where
88        P: StreamTextParser,
89    {
90        let mut all = StreamTextChunk::default();
91        for chunk in chunks {
92            let next = parser.push_str(chunk);
93            all.visible_text.push_str(&next.visible_text);
94            all.extracted.extend(next.extracted);
95        }
96        let tail = parser.finish();
97        all.visible_text.push_str(&tail.visible_text);
98        all.extracted.extend(tail.extracted);
99        all
100    }
101
102    #[test]
103    fn citation_parser_streams_across_chunk_boundaries() {
104        let mut parser = CitationStreamParser::new();
105        let out = collect_chunks(
106            &mut parser,
107            &[
108                "Hello <oai-mem-",
109                "citation>source A</oai-mem-",
110                "citation> world",
111            ],
112        );
113
114        assert_eq!(out.visible_text, "Hello  world");
115        assert_eq!(out.extracted, vec!["source A".to_string()]);
116    }
117
118    #[test]
119    fn citation_parser_buffers_partial_open_tag_prefix() {
120        let mut parser = CitationStreamParser::new();
121
122        let first = parser.push_str("abc <oai-mem-");
123        assert_eq!(first.visible_text, "abc ");
124        assert_eq!(first.extracted, Vec::<String>::new());
125
126        let second = parser.push_str("citation>x</oai-mem-citation>z");
127        let tail = parser.finish();
128
129        assert_eq!(second.visible_text, "z");
130        assert_eq!(second.extracted, vec!["x".to_string()]);
131        assert!(tail.is_empty());
132    }
133
134    #[test]
135    fn citation_parser_auto_closes_unterminated_tag_on_finish() {
136        let mut parser = CitationStreamParser::new();
137        let out = collect_chunks(&mut parser, &["x<oai-mem-citation>source"]);
138
139        assert_eq!(out.visible_text, "x");
140        assert_eq!(out.extracted, vec!["source".to_string()]);
141    }
142
143    #[test]
144    fn citation_parser_preserves_partial_open_tag_at_eof_if_not_a_full_tag() {
145        let mut parser = CitationStreamParser::new();
146        let out = collect_chunks(&mut parser, &["hello <oai-mem-"]);
147
148        assert_eq!(out.visible_text, "hello <oai-mem-");
149        assert_eq!(out.extracted, Vec::<String>::new());
150    }
151
152    #[test]
153    fn strip_citations_collects_all_citations() {
154        let (visible, citations) = strip_citations(
155            "a<oai-mem-citation>one</oai-mem-citation>b<oai-mem-citation>two</oai-mem-citation>c",
156        );
157
158        assert_eq!(visible, "abc");
159        assert_eq!(citations, vec!["one".to_string(), "two".to_string()]);
160    }
161
162    #[test]
163    fn strip_citations_auto_closes_unterminated_citation_at_eof() {
164        let (visible, citations) = strip_citations("x<oai-mem-citation>y");
165
166        assert_eq!(visible, "x");
167        assert_eq!(citations, vec!["y".to_string()]);
168    }
169
170    #[test]
171    fn citation_parser_does_not_support_nested_tags() {
172        let (visible, citations) = strip_citations(
173            "a<oai-mem-citation>x<oai-mem-citation>y</oai-mem-citation>z</oai-mem-citation>b",
174        );
175
176        assert_eq!(visible, "az</oai-mem-citation>b");
177        assert_eq!(citations, vec!["x<oai-mem-citation>y".to_string()]);
178    }
179}
codex_utils_stream_parser/citation.rs

codex_utils_stream_parser/
citation.rs
